jararaca 0.4.0a5__py3-none-any.whl → 0.4.0a19__py3-none-any.whl

jararaca/messagebus/worker.py

@@ -13,15 +13,7 @@ from abc import ABC
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
-from typing import (
-    Any,
-    AsyncContextManager,
-    AsyncGenerator,
-    Awaitable,
-    Optional,
-    Type,
-    get_origin,
-)
+from typing import Any, AsyncContextManager, AsyncGenerator, Awaitable, Optional, Type
 from urllib.parse import parse_qs, urlparse
 
 import aio_pika
@@ -35,7 +27,7 @@ from aio_pika.exceptions import (
     ChannelNotFoundEntity,
     ConnectionClosed,
 )
-from pydantic import BaseModel
+from pydantic import ValidationError
 
 from jararaca.broker_backend import MessageBrokerBackend
 from jararaca.broker_backend.mapper import get_message_broker_backend_from_url
@@ -65,9 +57,10 @@ from jararaca.microservice import (
     provide_shutdown_state,
     providing_app_type,
 )
+from jararaca.observability.hooks import record_exception, set_span_status
 from jararaca.scheduler.decorators import ScheduledActionData
 from jararaca.utils.rabbitmq_utils import RabbitmqUtils
-from jararaca.utils.retry import RetryConfig, retry_with_backoff
+from jararaca.utils.retry import RetryPolicy, retry_with_backoff
 
 logger = logging.getLogger(__name__)
 
@@ -77,16 +70,16 @@ class AioPikaWorkerConfig:
     url: str
     exchange: str
     prefetch_count: int
-    connection_retry_config: RetryConfig = field(
-        default_factory=lambda: RetryConfig(
+    connection_retry_config: RetryPolicy = field(
+        default_factory=lambda: RetryPolicy(
             max_retries=15,
             initial_delay=1.0,
             max_delay=60.0,
             backoff_factor=2.0,
         )
     )
-    consumer_retry_config: RetryConfig = field(
-        default_factory=lambda: RetryConfig(
+    consumer_retry_policy: RetryPolicy = field(
+        default_factory=lambda: RetryPolicy(
             max_retries=15,
             initial_delay=0.5,
             max_delay=40.0,
@@ -260,7 +253,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
             # Setup with retry
             await retry_with_backoff(
                 setup_consumer,
-                retry_config=self.config.consumer_retry_config,
+                retry_policy=self.config.consumer_retry_policy,
                 retry_exceptions=(
                     ChannelNotFoundEntity,
                     ChannelClosed,
@@ -321,7 +314,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
             # Setup with retry
             await retry_with_backoff(
                 setup_consumer,
-                retry_config=self.config.consumer_retry_config,
+                retry_policy=self.config.consumer_retry_policy,
                 retry_exceptions=(
                     ChannelNotFoundEntity,
                     ChannelClosed,
@@ -355,7 +348,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                 # Verify infrastructure with retry
                 infra_check_success = await retry_with_backoff(
                     self._verify_infrastructure,
-                    retry_config=self.config.connection_retry_config,
+                    retry_policy=self.config.connection_retry_config,
                     retry_exceptions=(Exception,),
                 )
 
@@ -503,12 +496,12 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                     len(pending_tasks),
                     ", ".join((task.get_name()) for task in pending_tasks),
                 )
-            else:
-                logger.warning("All in-flight tasks have completed.")
-            # Log any exceptions that occurred
-            # for result in results:
-            #     if isinstance(result, Exception):
-            #         logger.error("Task raised an exception during shutdown: %s", result)
+
+        logger.warning("All in-flight tasks have completed.")
+        # Log any exceptions that occurred
+        # for result in results:
+        #     if isinstance(result, Exception):
+        #         logger.error("Task raised an exception during shutdown: %s", result)
 
     async def close_channels_and_connection(self) -> None:
         """Close all channels and then the connection"""
@@ -609,7 +602,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
             # Create a new channel with retry
             channel = await retry_with_backoff(
                 fn=lambda: self._establish_channel(queue_name),
-                retry_config=self.config.consumer_retry_config,
+                retry_policy=self.config.consumer_retry_policy,
                 retry_exceptions=(
                     AMQPConnectionError,
                     AMQPChannelError,
@@ -667,7 +660,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
             # Create a new connection with retry
             self.connection = await retry_with_backoff(
                 self._establish_connection,
-                retry_config=self.config.connection_retry_config,
+                retry_policy=self.config.connection_retry_config,
                 retry_exceptions=(
                     AMQPConnectionError,
                     ConnectionError,
@@ -744,7 +737,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                     logger.warning(
                         "Error getting channel for %s, retrying: %s", queue_name, e
                     )
-                    await asyncio.sleep(retry_delay)
+                    await self._wait_delay_or_shutdown(retry_delay)
                     retry_delay *= 2
                 else:
                     logger.error(
@@ -755,6 +748,24 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                     )
                     raise
 
+    async def _wait_delay_or_shutdown(self, delay: float) -> None:
+        """
+        Wait for the specified delay or exit early if shutdown is initiated.
+
+        Args:
+            delay: Delay in seconds to wait
+        """
+
+        wait_cor = asyncio.create_task(asyncio.sleep(delay), name="delayed-retry-wait")
+        wait_shutdown_cor = asyncio.create_task(
+            self.shutdown_event.wait(), name="delayed-retry-shutdown-wait"
+        )
+
+        await asyncio.wait(
+            [wait_cor, wait_shutdown_cor],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
     async def _monitor_connection_health(self) -> None:
         """
         Monitor connection health and trigger shutdown if connection is lost.
@@ -762,7 +773,9 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         """
         while not self.shutdown_event.is_set():
             try:
-                await asyncio.sleep(self.config.connection_health_check_interval)
+                await self._wait_delay_or_shutdown(
+                    self.config.connection_health_check_interval
+                )
 
                 if self.shutdown_event.is_set():
                     break
@@ -780,7 +793,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                 break
             except Exception as e:
                 logger.error("Error in connection health monitoring: %s", e)
-                await asyncio.sleep(5)  # Wait before retrying
+                await self._wait_delay_or_shutdown(5)  # Wait before retrying
 
     async def _is_connection_healthy(self) -> bool:
         """
@@ -890,9 +903,9 @@ def create_message_bus(
         prefetch_count = int(query_params["prefetch_count"][0])
 
         # Parse optional retry configuration parameters
-        connection_retry_config = RetryConfig()
-        consumer_retry_config = RetryConfig(
-            max_retries=30, initial_delay=5, max_delay=60.0, backoff_factor=3.0
+        connection_retry_config = RetryPolicy()
+        consumer_retry_config = RetryPolicy(
+            max_retries=5, initial_delay=5, max_delay=60.0, backoff_factor=3.0
         )
 
         # Parse heartbeat and health check intervals
@@ -987,7 +1000,7 @@ def create_message_bus(
             exchange=exchange,
             prefetch_count=prefetch_count,
             connection_retry_config=connection_retry_config,
-            consumer_retry_config=consumer_retry_config,
+            consumer_retry_policy=consumer_retry_config,
             connection_heartbeat_interval=connection_heartbeat_interval,
             connection_health_check_interval=connection_health_check_interval,
         )
@@ -1049,8 +1062,8 @@ class ScheduledMessageHandlerCallback:
             try:
                 if not self.consumer.connection_healthy:
                     # Still not healthy, requeue the message
-                    async with self.consumer.get_channel_ctx(self.queue_name):
-                        await aio_pika_message.reject(requeue=True)
+
+                    await aio_pika_message.reject(requeue=True)
                     return
             except Exception as e:
                 logger.error(
@@ -1088,8 +1101,8 @@ class ScheduledMessageHandlerCallback:
             )
             try:
                 # Use channel context for requeuing
-                async with self.consumer.get_channel_ctx(self.queue_name):
-                    await aio_pika_message.reject(requeue=True)
+
+                await aio_pika_message.reject(requeue=True)
                 return
             except RuntimeError:
                 logger.warning(
@@ -1106,8 +1119,8 @@ class ScheduledMessageHandlerCallback:
                 self.queue_name,
             )
             try:
-                async with self.consumer.get_channel_ctx(self.queue_name):
-                    await aio_pika_message.reject(requeue=True)
+
+                await aio_pika_message.reject(requeue=True)
                 return
             except Exception as e:
                 logger.error(
@@ -1203,8 +1216,8 @@ class MessageHandlerCallback:
             )
             try:
                 # Use channel context for requeuing
-                async with self.consumer.get_channel_ctx(self.queue_name):
-                    await aio_pika_message.reject(requeue=True)
+
+                await aio_pika_message.reject(requeue=True)
             except RuntimeError:
                 logger.warning(
                     "Could not requeue message during shutdown - channel not available"
@@ -1221,8 +1234,8 @@ class MessageHandlerCallback:
             try:
                 if not self.consumer.connection_healthy:
                     # Still not healthy, requeue the message
-                    async with self.consumer.get_channel_ctx(self.queue_name):
-                        await aio_pika_message.reject(requeue=True)
+
+                    await aio_pika_message.reject(requeue=True)
                     return
             except Exception as e:
                 logger.error(
@@ -1236,18 +1249,22 @@ class MessageHandlerCallback:
                 name=f"MessageHandler-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
             )
             self.consumer.tasks.add(task)
-            task.add_done_callback(self.handle_message_consume_done)
 
-    def handle_message_consume_done(self, task: asyncio.Task[Any]) -> None:
-        self.consumer.tasks.discard(task)
-        if task.cancelled():
-            logger.warning("Task for queue %s was cancelled", self.queue_name)
-            return
+            def handle_message_consume_done(task: asyncio.Task[Any]) -> None:
+                self.consumer.tasks.discard(task)
+                if task.cancelled():
+                    logger.warning("Task for queue %s was cancelled", self.queue_name)
+                    return
 
-        if (error := task.exception()) is not None:
-            logger.exception(
-                "Error processing message for queue %s", self.queue_name, exc_info=error
-            )
+                if (error := task.exception()) is not None:
+                    logger.exception(
+                        "Error processing message id %s for queue %s",
+                        aio_pika_message.message_id,
+                        self.queue_name,
+                        exc_info=error,
+                    )
+
+            task.add_done_callback(handle_message_consume_done)
 
     async def __call__(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
@@ -1257,6 +1274,8 @@ class MessageHandlerCallback:
     async def handle_reject_message(
         self,
         aio_pika_message: aio_pika.abc.AbstractIncomingMessage,
+        *,
+        requeue_timeout: float = 0,
         requeue: bool = False,
         retry_count: int = 0,
         exception: Optional[BaseException] = None,
@@ -1280,7 +1299,10 @@ class MessageHandlerCallback:
                 and exception is not None
             ):
                 # Get retry config from consumer
-                retry_config = self.consumer.config.consumer_retry_config
+                retry_config = (
+                    self.message_handler.spec.retry_config
+                    or self.consumer.config.consumer_retry_policy
+                )
 
                 # Check if we reached max retries
                 if retry_count >= retry_config.max_retries:
@@ -1293,7 +1315,6 @@ class MessageHandlerCallback:
                     )
                     # Dead-letter the message after max retries
                     try:
-
                         await aio_pika_message.reject(requeue=False)
                     except Exception as e:
                         logger.error(
@@ -1332,12 +1353,13 @@ class MessageHandlerCallback:
                 }
 
                 # Schedule retry after delay
-                asyncio.create_task(
+                task = asyncio.create_task(
                     self._delayed_retry(
                         aio_pika_message, delay, retry_count + 1, exception
                     ),
                     name=f"MessageHandler-{self.queue_name}-delayed-retry-{message_id}",
                 )
+                self.consumer.tasks.add(task)
 
                 # Acknowledge the current message since we'll handle retry ourselves
                 try:
@@ -1350,20 +1372,22 @@ class MessageHandlerCallback:
 
             # Standard reject without retry or with immediate requeue
             try:
-                async with self.consumer.get_channel_ctx(self.queue_name):
-                    await aio_pika_message.reject(requeue=requeue)
-                    if requeue:
-                        logger.warning(
-                            "Message %s (%s) requeued for immediate retry",
-                            message_id,
-                            self.queue_name,
-                        )
-                    else:
-                        logger.warning(
-                            "Message %s (%s) rejected without requeue",
-                            message_id,
-                            self.queue_name,
-                        )
+                await self._wait_delay_or_shutdown(
+                    requeue_timeout
+                )  # Optional delay before requeueing
+                await aio_pika_message.reject(requeue=requeue)
+                if requeue:
+                    logger.warning(
+                        "Message %s (%s) requeued for immediate retry",
+                        message_id,
+                        self.queue_name,
+                    )
+                else:
+                    logger.warning(
+                        "Message %s (%s) rejected without requeue",
+                        message_id,
+                        self.queue_name,
+                    )
             except Exception as e:
                 logger.error("Failed to reject message %s: %s", message_id, e)
 
@@ -1452,7 +1476,9 @@ class MessageHandlerCallback:
                             attempt + 1,
                             e,
                         )
-                        await asyncio.sleep(1.0 * (attempt + 1))  # Exponential backoff
+                        await self._wait_delay_or_shutdown(
+                            (1.0 * (attempt + 1))
+                        )  # Exponential backoff
                     else:
                         logger.error(
                             "Failed to republish message %s after %s attempts: %s",
@@ -1477,22 +1503,7 @@ class MessageHandlerCallback:
                 pass
 
     async def _wait_delay_or_shutdown(self, delay: float) -> None:
-        """
-        Wait for the specified delay or exit early if shutdown is initiated.
-
-        Args:
-            delay: Delay in seconds to wait
-        """
-
-        wait_cor = asyncio.create_task(asyncio.sleep(delay), name="delayed-retry-wait")
-        wait_shutdown_cor = asyncio.create_task(
-            self.consumer.shutdown_event.wait(), name="delayed-retry-shutdown-wait"
-        )
-
-        await asyncio.wait(
-            [wait_cor, wait_shutdown_cor],
-            return_when=asyncio.FIRST_COMPLETED,
-        )
+        await self.consumer._wait_delay_or_shutdown(delay)
 
     async def handle_message(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
@@ -1508,47 +1519,53 @@ class MessageHandlerCallback:
         handler_data = self.message_handler
 
         handler = handler_data.instance_callable
+        handler_method = handler_data.controller_member.member_function
 
-        sig = inspect.signature(handler)
+        # sig = inspect.signature(handler)
 
-        if len(sig.parameters) != 1:
-            logger.warning(
-                "Handler for topic '%s' must have exactly one parameter which is MessageOf[T extends Message]"
-                % routing_key
-            )
-            return
+        # if len(sig.parameters) != 1:
+        #     logger.warning(
+        #         "Handler for topic '%s' must have exactly one parameter which is MessageOf[T extends Message]"
+        #         % routing_key
+        #     )
+        #     return
 
-        parameter = list(sig.parameters.values())[0]
+        # parameter = list(sig.parameters.values())[0]
 
-        param_origin = get_origin(parameter.annotation)
+        # param_origin = get_origin(parameter.annotation)
 
-        if param_origin is not MessageOf:
-            logger.warning(
-                "Handler for topic '%s' must have exactly one parameter of type Message"
-                % routing_key
-            )
-            return
+        # if param_origin is not MessageOf:
+        #     logger.warning(
+        #         "Handler for topic '%s' must have exactly one parameter of type Message"
+        #         % routing_key
+        #     )
+        #     return
 
-        if len(parameter.annotation.__args__) != 1:
-            logger.warning(
-                "Handler for topic '%s' must have exactly one parameter of type Message"
-                % routing_key
-            )
-            return
+        # if len(parameter.annotation.__args__) != 1:
+        #     logger.warning(
+        #         "Handler for topic '%s' must have exactly one parameter of type Message"
+        #         % routing_key
+        #     )
+        #     return
 
-        message_type = parameter.annotation.__args__[0]
+        # message_type = parameter.annotation.__args__[0]
 
-        if not issubclass(message_type, BaseModel):
-            logger.warning(
-                "Handler for topic '%s' must have exactly one parameter of type MessageOf[BaseModel]"
-                % routing_key
-            )
-            return
+        # if not issubclass(message_type, BaseModel):
+        #     logger.warning(
+        #         "Handler for topic '%s' must have exactly one parameter of type MessageOf[BaseModel]"
+        #         % routing_key
+        #     )
+        #     return
 
-        builded_message = AioPikaMessage(aio_pika_message, message_type)
+        mode, message_type = MessageHandler.validate_decorated_fn(handler_method)
+
+        built_message = AioPikaMessage(aio_pika_message, message_type)
 
         incoming_message_spec = MessageHandler.get_last(handler)
-        assert incoming_message_spec is not None
+        assert incoming_message_spec is not None, "Incoming message spec must be set"
+        # Extract retry count from headers if available
+        headers = aio_pika_message.headers or {}
+        retry_count = int(str(headers.get("x-retry-count", 0)))
 
         with provide_implicit_headers(aio_pika_message.headers), provide_shutdown_state(
             self.consumer.shutdown_state
@@ -1557,8 +1574,10 @@ class MessageHandlerCallback:
                 AppTransactionContext(
                     controller_member_reflect=handler_data.controller_member,
                     transaction_data=MessageBusTransactionData(
+                        message_id=aio_pika_message.message_id,
+                        processing_attempt=retry_count + 1,
                         message_type=message_type,
-                        message=builded_message,
+                        message=built_message,
                         topic=routing_key,
                     ),
                 )
@@ -1575,94 +1594,149 @@ class MessageHandlerCallback:
                         with provide_bus_message_controller(
                             AioPikaMessageBusController(aio_pika_message)
                         ):
-                            await handler(builded_message)
-                            with suppress(aio_pika.MessageProcessError):
-                                # Use channel context for acknowledgement with retry
-                                try:
-                                    await aio_pika_message.ack()
-                                except Exception as ack_error:
-                                    logger.warning(
-                                        "Failed to acknowledge message %s: %s",
-                                        aio_pika_message.message_id or "unknown",
-                                        ack_error,
-                                    )
-                        successfully = True
-                    except BaseException as base_exc:
-                        successfully = False
-                        # Get message id for logging
-                        message_id = aio_pika_message.message_id or "unknown"
-
-                        # Extract retry count from headers if available
-                        headers = aio_pika_message.headers or {}
-                        retry_count = int(str(headers.get("x-retry-count", 0)))
-
-                        # Process exception handler if configured
-                        if incoming_message_spec.exception_handler is not None:
                             try:
-                                incoming_message_spec.exception_handler(base_exc)
-                            except Exception as nested_exc:
-                                logger.exception(
-                                    "Error processing exception handler for message %s: %s | %s",
-                                    message_id,
+                                if mode == "WRAPPED":
+                                    future = handler(built_message)
+                                else:
+                                    try:
+
+                                        payload = built_message.payload()
+                                    except ValidationError as exc:
+                                        logger.exception(
+                                            "Validation error parsing message %s on topic %s",
+                                            aio_pika_message.message_id or "unknown",
+                                            routing_key,
+                                        )
+                                        aio_pika_message.headers["x-last-error"] = (
+                                            "Validation error parsing message payload"
+                                        )
+                                        await aio_pika_message.reject(requeue=False)
+                                        record_exception(
+                                            exc,
+                                        )
+                                        set_span_status("ERROR")
+                                        return
+                                    future = handler(payload)
+
+                                await future
+
+                                with suppress(aio_pika.MessageProcessError):
+                                    # Use channel context for acknowledgement with retry
+                                    try:
+                                        await aio_pika_message.ack()
+                                        set_span_status("OK")
+                                    except Exception as ack_error:
+                                        logger.warning(
+                                            "Failed to acknowledge message %s: %s",
+                                            aio_pika_message.message_id or "unknown",
+                                            ack_error,
+                                        )
+                                successfully = True
+                            except Exception as base_exc:
+                                set_span_status("ERROR")
+                                record_exception(
                                     base_exc,
-                                    nested_exc,
+                                    {
+                                        "message_id": aio_pika_message.message_id
+                                        or "unknown",
+                                        "routing_key": routing_key,
+                                    },
                                 )
-                        else:
-                            logger.exception(
-                                "Error processing message %s on topic %s: %s",
-                                message_id,
-                                routing_key,
-                                str(base_exc),
-                            )
+                                successfully = False
+                                # Get message id for logging
+                                message_id = aio_pika_message.message_id or "unknown"
+
+                                # Process exception handler if configured
+                                if incoming_message_spec.exception_handler is not None:
+                                    try:
+                                        incoming_message_spec.exception_handler(
+                                            base_exc
+                                        )
+                                    except Exception as nested_exc:
+                                        logger.exception(
+                                            "Error processing exception handler for message %s: %s | %s",
+                                            message_id,
+                                            base_exc,
+                                            nested_exc,
+                                        )
+                                else:
+                                    logger.exception(
+                                        "Error processing message %s on topic %s: %s",
+                                        message_id,
+                                        routing_key,
+                                        str(base_exc),
+                                    )
 
-                        # Handle rejection with retry logic
-                        if incoming_message_spec.nack_on_exception:
-                            await self.handle_reject_message(
-                                aio_pika_message,
-                                requeue=False,  # Don't requeue directly, use our backoff mechanism
-                                retry_count=retry_count,
-                                exception=base_exc,
+                                # Handle rejection with retry logic
+                                if incoming_message_spec.nack_on_exception:
+                                    await self.handle_reject_message(
+                                        aio_pika_message,
+                                        requeue_timeout=incoming_message_spec.nack_delay_on_exception,
+                                        requeue=False,  # Don't requeue directly, use our backoff mechanism
+                                        retry_count=retry_count,
+                                        exception=base_exc,
+                                    )
+                                else:
+                                    # Message shouldn't be retried, reject it
+                                    await self.handle_reject_message(
+                                        aio_pika_message,
+                                        requeue=False,
+                                        requeue_timeout=incoming_message_spec.nack_delay_on_exception,
+                                        exception=base_exc,
+                                    )
+
+                            elapsed_time = time.perf_counter() - start_time
+                            # Message processed successfully, log and clean up any retry state
+                            message_id = aio_pika_message.message_id or str(
+                                uuid.uuid4()
                             )
-                        else:
-                            # Message shouldn't be retried, reject it
-                            await self.handle_reject_message(
-                                aio_pika_message, requeue=False, exception=base_exc
+                            if message_id in self.retry_state:
+                                del self.retry_state[message_id]
+
+                            # Log success with retry information if applicable
+                            headers = aio_pika_message.headers or {}
+                            traceparent = headers.get("traceparent")
+                            trace_info = (
+                                f" [traceparent={str(traceparent)}]"
+                                if traceparent
+                                else ""
                             )
 
-                    elapsed_time = time.perf_counter() - start_time
-                    # Message processed successfully, log and clean up any retry state
-                    message_id = aio_pika_message.message_id or str(uuid.uuid4())
-                    if message_id in self.retry_state:
-                        del self.retry_state[message_id]
-
-                    # Log success with retry information if applicable
-                    headers = aio_pika_message.headers or {}
-                    traceparent = headers.get("traceparent")
-                    trace_info = (
-                        f" [traceparent={str(traceparent)}]" if traceparent else ""
-                    )
+                            if "x-retry-count" in headers:
+                                retry_count = int(str(headers.get("x-retry-count", 0)))
+                                logger.debug(
+                                    "Message %s#%s processed "
+                                    + (
+                                        "successfully"
+                                        if successfully
+                                        else "with errors"
+                                    )
+                                    + " after %s retries in %.4fs%s",
+                                    message_id,
+                                    self.queue_name,
+                                    retry_count,
+                                    elapsed_time,
+                                    trace_info,
+                                )
+                            else:
+                                logger.debug(
+                                    "Message %s#%s processed "
+                                    + (
+                                        "successfully"
+                                        if successfully
+                                        else "with errors"
+                                    )
+                                    + " in %.4fs%s",
+                                    message_id,
+                                    self.queue_name,
+                                    elapsed_time,
+                                    trace_info,
+                                )
+                                ...
 
-                    if "x-retry-count" in headers:
-                        retry_count = int(str(headers.get("x-retry-count", 0)))
-                        logger.debug(
-                            "Message %s#%s processed "
-                            + ("successfully" if successfully else "with errors")
-                            + " after %s retries in %.4fs%s",
-                            message_id,
-                            self.queue_name,
-                            retry_count,
-                            elapsed_time,
-                            trace_info,
-                        )
-                    else:
-                        logger.debug(
-                            "Message %s#%s processed "
-                            + ("successfully" if successfully else "with errors")
-                            + " in %.4fs%s",
-                            message_id,
-                            self.queue_name,
-                            elapsed_time,
-                            trace_info,
+                    except Exception as base_exc:
+                        logger.critical(
+                            f"Critical error processing message {aio_pika_message.message_id} when providing bus message controller: {base_exc}"
                         )
 
 
@@ -1841,13 +1915,15 @@ class AioPikaMessageBusController(BusMessageController):
     async def reject(self) -> None:
         await self.aio_pika_message.reject()
 
-    async def retry(self) -> None:
+    async def retry(self, delay: float = 5) -> None:
         """
         Retry the message immediately by rejecting with requeue flag.
         This doesn't use the exponential backoff mechanism.
         """
         callback = self._get_callback()
-        await callback.handle_reject_message(self.aio_pika_message, requeue=True)
+        await callback.handle_reject_message(
+            self.aio_pika_message, requeue=True, requeue_timeout=delay
+        )
 
     async def retry_later(self, delay: int) -> None:
         """