jararaca 0.3.11a16__py3-none-any.whl → 0.4.0a19__py3-none-any.whl

jararaca/messagebus/worker.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2025 Lucas S
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 import asyncio
 import inspect
 import logging
@@ -9,22 +13,21 @@ 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
 import aio_pika.abc
 import uvloop
-from aio_pika.exceptions import AMQPError, ChannelClosed, ChannelNotFoundEntity
-from pydantic import BaseModel
+from aio_pika.exceptions import (
+    AMQPChannelError,
+    AMQPConnectionError,
+    AMQPError,
+    ChannelClosed,
+    ChannelNotFoundEntity,
+    ConnectionClosed,
+)
+from pydantic import ValidationError
 
 from jararaca.broker_backend import MessageBrokerBackend
 from jararaca.broker_backend.mapper import get_message_broker_backend_from_url
@@ -43,16 +46,21 @@ from jararaca.messagebus.decorators import (
     MessageHandlerData,
     ScheduleDispatchData,
 )
+from jararaca.messagebus.implicit_headers import provide_implicit_headers
 from jararaca.messagebus.message import Message, MessageOf
 from jararaca.microservice import (
     AppTransactionContext,
     MessageBusTransactionData,
     Microservice,
     SchedulerTransactionData,
+    ShutdownState,
+    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__)
 
@@ -62,22 +70,25 @@ 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,
             backoff_factor=2.0,
         )
     )
+    # Connection health monitoring settings
+    connection_heartbeat_interval: float = 30.0  # seconds
+    connection_health_check_interval: float = 10.0  # seconds
 
 
 class AioPikaMessage(MessageOf[Message]):
@@ -129,6 +140,20 @@ class MessageBusConsumer(ABC):
         """Close all resources related to the consumer"""
 
 
+class _WorkerShutdownState(ShutdownState):
+    def __init__(self, shutdown_event: asyncio.Event):
+        self.shutdown_event = shutdown_event
+
+    def request_shutdown(self) -> None:
+        self.shutdown_event.set()
+
+    def is_shutdown_requested(self) -> bool:
+        return self.shutdown_event.is_set()
+
+    async def wait_for_shutdown(self) -> None:
+        await self.shutdown_event.wait()
+
+
 class AioPikaMicroserviceConsumer(MessageBusConsumer):
     def __init__(
         self,
@@ -146,11 +171,18 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         self.incoming_map: dict[str, MessageHandlerData] = {}
         self.uow_context_provider = uow_context_provider
         self.shutdown_event = asyncio.Event()
+        self.shutdown_state = _WorkerShutdownState(self.shutdown_event)
         self.lock = asyncio.Lock()
         self.tasks: set[asyncio.Task[Any]] = set()
         self.connection: aio_pika.abc.AbstractConnection | None = None
         self.channels: dict[str, aio_pika.abc.AbstractChannel] = {}
 
+        # Connection resilience attributes
+        self.connection_healthy = False
+        self.connection_lock = asyncio.Lock()
+        self.consumer_tags: dict[str, str] = {}  # Track consumer tags for cleanup
+        self.health_check_task: asyncio.Task[Any] | None = None
+
     async def _verify_infrastructure(self) -> bool:
         """
         Verify that the required RabbitMQ infrastructure (exchanges, queues) exists.
@@ -188,31 +220,48 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         async def setup_consumer() -> None:
             # Create a channel using the context manager
             async with self.create_channel(queue_name) as channel:
-                queue = await RabbitmqUtils.get_queue(
+                queue: aio_pika.abc.AbstractQueue = await RabbitmqUtils.get_queue(
                     channel=channel, queue_name=queue_name
                 )
 
-                # Configure consumer right away while in the context
-                await queue.consume(
+                # Configure consumer and get the consumer tag
+                consumer_tag = await queue.consume(
                     callback=MessageHandlerCallback(
                         consumer=self,
                         queue_name=queue_name,
                         routing_key=routing_key,
                         message_handler=handler,
                     ),
-                    no_ack=handler.spec.auto_ack,
+                    # no_ack=handler.spec.auto_ack,
                 )
 
+                # Store consumer tag for cleanup
+                self.consumer_tags[queue_name] = consumer_tag
+
                 logger.info(
-                    f"Consuming message handler {queue_name} on dedicated channel"
+                    "Consuming message handler %s on dedicated channel", queue_name
+                )
+
+                await self.shutdown_event.wait()
+
+                logger.warning(
+                    "Shutdown event received, stopping consumer for %s", queue_name
                 )
+                await queue.cancel(consumer_tag)
 
         try:
             # Setup with retry
             await retry_with_backoff(
                 setup_consumer,
-                retry_config=self.config.consumer_retry_config,
-                retry_exceptions=(ChannelNotFoundEntity, ChannelClosed, AMQPError),
+                retry_policy=self.config.consumer_retry_policy,
+                retry_exceptions=(
+                    ChannelNotFoundEntity,
+                    ChannelClosed,
+                    AMQPError,
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ConnectionClosed,
+                ),
             )
             return True
         except Exception as e:
@@ -238,8 +287,8 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                     channel=channel, queue_name=queue_name
                 )
 
-                # Configure consumer right away while in the context
-                await queue.consume(
+                # Configure consumer and get the consumer tag
+                consumer_tag = await queue.consume(
                     callback=ScheduledMessageHandlerCallback(
                         consumer=self,
                         queue_name=queue_name,
@@ -249,14 +298,31 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                     no_ack=True,
                 )
 
-                logger.info(f"Consuming scheduler {queue_name} on dedicated channel")
+                # Store consumer tag for cleanup
+                self.consumer_tags[queue_name] = consumer_tag
+
+                logger.debug("Consuming scheduler %s on dedicated channel", queue_name)
+
+                await self.shutdown_event.wait()
+
+                logger.warning(
+                    "Shutdown event received, stopping consumer for %s", queue_name
+                )
+                await queue.cancel(consumer_tag)
 
         try:
             # Setup with retry
             await retry_with_backoff(
                 setup_consumer,
-                retry_config=self.config.consumer_retry_config,
-                retry_exceptions=(ChannelNotFoundEntity, ChannelClosed, AMQPError),
+                retry_policy=self.config.consumer_retry_policy,
+                retry_exceptions=(
+                    ChannelNotFoundEntity,
+                    ChannelClosed,
+                    AMQPError,
+                    AMQPConnectionError,
+                    AMQPChannelError,
+                    ConnectionClosed,
+                ),
             )
             return True
         except Exception as e:
@@ -269,160 +335,195 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         """
         Main consume method that sets up all message handlers and scheduled actions with retry mechanisms.
         """
-        # Verify infrastructure with retry
-        infra_check_success = await retry_with_backoff(
-            self._verify_infrastructure,
-            retry_config=self.config.connection_retry_config,
-            retry_exceptions=(Exception,),
-        )
-
-        if not infra_check_success:
-            logger.critical("Failed to verify RabbitMQ infrastructure. Shutting down.")
-            self.shutdown_event.set()
-            return
-
-        async def wait_for(
-            type: str, name: str, coroutine: Awaitable[bool]
-        ) -> tuple[str, str, bool]:
-            return type, name, await coroutine
+        # Establish initial connection
+        try:
+            async with self.connect() as connection:
+                self.connection_healthy = True
 
-        tasks: set[asyncio.Task[tuple[str, str, bool]]] = set()
+                # Start connection health monitoring
+                self.health_check_task = asyncio.create_task(
+                    self._monitor_connection_health(), name="ConnectionHealthMonitor"
+                )
 
-        # Setup message handlers
-        for handler in self.message_handler_set:
-            queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.instance_callable.__module__}.{handler.instance_callable.__qualname__}"
-            self.incoming_map[queue_name] = handler
+                # Verify infrastructure with retry
+                infra_check_success = await retry_with_backoff(
+                    self._verify_infrastructure,
+                    retry_policy=self.config.connection_retry_config,
+                    retry_exceptions=(Exception,),
+                )
 
-            tasks.add(
-                task := asyncio.create_task(
-                    wait_for(
-                        "message_handler",
-                        queue_name,
-                        self._setup_message_handler_consumer(handler),
+                if not infra_check_success:
+                    logger.critical(
+                        "Failed to verify RabbitMQ infrastructure. Shutting down."
                     )
-                )
-            )
-            # task.add_done_callback(tasks.discard)
-            # success = await self._setup_message_handler_consumer(handler)
-            # if not success:
-            #     logger.warning(
-            #         f"Failed to set up consumer for {queue_name}, will not process messages from this queue"
-            #     )
-
-        # Setup scheduled actions
-        for scheduled_action in self.scheduled_actions:
-
-            queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
-            tasks.add(
-                task := asyncio.create_task(
-                    wait_for(
-                        "scheduled_action",
-                        queue_name,
-                        self._setup_scheduled_action_consumer(scheduled_action),
+                    self.shutdown_event.set()
+                    return
+
+                async def wait_for(
+                    type: str, name: str, coroutine: Awaitable[bool]
+                ) -> tuple[str, str, bool]:
+                    return type, name, await coroutine
+
+                tasks: set[asyncio.Task[tuple[str, str, bool]]] = set()
+
+                # Setup message handlers
+                for handler in self.message_handler_set:
+                    queue_name = f"{handler.message_type.MESSAGE_TOPIC}.{handler.instance_callable.__module__}.{handler.instance_callable.__qualname__}"
+                    self.incoming_map[queue_name] = handler
+
+                    tasks.add(
+                        task := asyncio.create_task(
+                            wait_for(
+                                "message_handler",
+                                queue_name,
+                                self._setup_message_handler_consumer(handler),
+                            ),
+                            name=f"MessageHandler-{queue_name}-setup-consumer",
+                        )
                     )
-                )
-            )
-            # task.add_done_callback(tasks.discard)
-
-            # success = await self._setup_scheduled_action_consumer(scheduled_action)
-            # if not success:
-            #     queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
-            #     logger.warning(
-            #         f"Failed to set up consumer for scheduled action {queue_name}, will not process scheduled tasks from this queue"
-            #     )
-
-        async def handle_task_results() -> None:
-            for task in asyncio.as_completed(tasks):
-                type, name, success = await task
-                if success:
-                    logger.info(f"Successfully set up {type} consumer for {name}")
-                else:
-                    logger.warning(
-                        f"Failed to set up {type} consumer for {name}, will not process messages from this queue"
+
+                # Setup scheduled actions
+                for scheduled_action in self.scheduled_actions:
+                    queue_name = f"{scheduled_action.callable.__module__}.{scheduled_action.callable.__qualname__}"
+                    tasks.add(
+                        task := asyncio.create_task(
+                            wait_for(
+                                "scheduled_action",
+                                queue_name,
+                                self._setup_scheduled_action_consumer(scheduled_action),
+                            ),
+                            name=f"ScheduledAction-{queue_name}-setup-consumer",
+                        )
                     )
 
-        handle_task_results_task = asyncio.create_task(handle_task_results())
+                async def handle_task_results() -> None:
+                    for task in asyncio.as_completed(tasks):
+                        type, name, success = await task
+                        if success:
+                            logger.debug(
+                                "Successfully set up %s consumer for %s", type, name
+                            )
+                        else:
+                            logger.warning(
+                                "Failed to set up %s consumer for %s, will not process messages from this queue",
+                                type,
+                                name,
+                            )
 
-        # Wait for shutdown signal
-        await self.shutdown_event.wait()
-        logger.info("Shutdown event received, stopping consumers")
-        handle_task_results_task.cancel()
-        with suppress(asyncio.CancelledError):
-            await handle_task_results_task
-        for task in tasks:
-            if not task.done():
-                task.cancel()
-                with suppress(asyncio.CancelledError):
-                    await task
-        logger.info("Worker shutting down")
+                handle_task_results_task = asyncio.create_task(
+                    handle_task_results(), name="HandleSetupTaskResults"
+                )
 
-        # Wait for all tasks to complete
-        await self.wait_all_tasks_done()
+                # Wait for shutdown signal
+                await self.shutdown_event.wait()
+                logger.debug("Shutdown event received, stopping consumers")
 
-        # Close all channels and the connection
-        await self.close_channels_and_connection()
+                await self.cancel_queue_consumers()
 
-    async def wait_all_tasks_done(self) -> None:
-        if not self.tasks:
-            return
+                # Cancel health monitoring
+                if self.health_check_task:
+                    self.health_check_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await self.health_check_task
 
-        logger.info(f"Waiting for {len(self.tasks)} in-flight tasks to complete")
-        async with self.lock:
-            # Use gather with return_exceptions=True to ensure all tasks are awaited
-            # even if some raise exceptions
-            results = await asyncio.gather(*self.tasks, return_exceptions=True)
+                handle_task_results_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await handle_task_results_task
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                        with suppress(asyncio.CancelledError):
+                            await task
+                logger.debug("Worker shutting down")
+                # Wait for all tasks to complete
+                await self.wait_all_tasks_done()
+
+                # Close all channels and the connection
+                await self.close_channels_and_connection()
 
-            # Log any exceptions that occurred
-            for result in results:
-                if isinstance(result, Exception):
-                    logger.error(f"Task raised an exception during shutdown: {result}")
+        except Exception as e:
+            logger.critical("Failed to establish initial connection to RabbitMQ: %s", e)
+            # Re-raise the exception so it can be caught by the caller
+            raise
 
-    async def close_channels_and_connection(self) -> None:
-        """Close all channels and then the connection"""
-        # Close all channels
-        channel_close_tasks = []
+    async def cancel_queue_consumers(self) -> None:
+        """
+        Cancel all active queue consumers.
+        """
+        logger.debug("Cancelling all active queue consumers...")
         for queue_name, channel in self.channels.items():
             try:
                 if not channel.is_closed:
-                    logger.info(f"Closing channel for queue {queue_name}")
-                    channel_close_tasks.append(channel.close())
-                else:
-                    logger.info(f"Channel for queue {queue_name} already closed")
+                    # Cancel consumer if we have its tag
+                    if queue_name in self.consumer_tags:
+                        try:
+                            queue = await channel.get_queue(queue_name, ensure=False)
+                            if queue:
+                                await queue.cancel(self.consumer_tags[queue_name])
+                        except Exception as cancel_error:
+                            logger.warning(
+                                "Error cancelling consumer for %s: %s",
+                                queue_name,
+                                cancel_error,
+                            )
+                        del self.consumer_tags[queue_name]
             except Exception as e:
-                logger.error(
-                    f"Error preparing to close channel for queue {queue_name}: {e}"
-                )
+                logger.warning("Error cancelling consumer for %s: %s", queue_name, e)
 
-        # Wait for all channels to close (if any)
-        if channel_close_tasks:
-            try:
-                await asyncio.gather(*channel_close_tasks, return_exceptions=True)
-            except Exception as e:
-                logger.error(f"Error during channel closures: {e}")
+    async def wait_all_tasks_done(self) -> None:
+        if not self.tasks:
+            return
 
-        # Clear channels dictionary
-        self.channels.clear()
+        logger.warning(
+            "Waiting for (%s) in-flight tasks to complete: %s",
+            len(self.tasks),
+            ", ".join((task.get_name()) for task in self.tasks),
+        )
+        # async with self.lock:
+        # Use gather with return_exceptions=True to ensure all tasks are awaited
+        # even if some raise exceptions
+        # results = await asyncio.gather(*self.tasks, return_exceptions=True)
+        pending_tasks = [task for task in self.tasks if not task.done()]
+        while len(pending_tasks) > 0:
+            if not pending_tasks:
+                break
+            await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            pending_tasks = [task for task in pending_tasks if not task.done()]
+            if len(pending_tasks) > 0:
+                logger.warning(
+                    "Waiting for (%s) in-flight tasks to complete: %s",
+                    len(pending_tasks),
+                    ", ".join((task.get_name()) for task in pending_tasks),
+                )
 
-        # Close the connection
-        if self.connection:
-            try:
-                if not self.connection.is_closed:
-                    logger.info("Closing RabbitMQ connection")
-                    await self.connection.close()
-                else:
-                    logger.info("RabbitMQ connection already closed")
-            except Exception as e:
-                logger.error(f"Error closing RabbitMQ connection: {e}")
-            self.connection = None
+        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"""
+        logger.warning("Closing channels and connection...")
+        await self._cleanup_connection()
 
     def shutdown(self) -> None:
         """Signal for shutdown"""
-        logger.info("Initiating graceful shutdown")
+        logger.warning("Initiating graceful shutdown")
         self.shutdown_event.set()
 
     async def close(self) -> None:
         """Implement MessageBusConsumer.close for cleanup"""
+        logger.warning("Closing consumer...")
         self.shutdown()
+
+        # Cancel health monitoring
+        if self.health_check_task:
+            self.health_check_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.health_check_task
+
         await self.wait_all_tasks_done()
         await self.close_channels_and_connection()
 
@@ -432,25 +533,45 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         This helps with error handling when a channel might have been closed.
         """
         if queue_name not in self.channels:
-            logger.warning(f"No channel found for queue {queue_name}")
+            logger.warning("No channel found for queue %s", queue_name)
             return None
 
         try:
             channel = self.channels[queue_name]
             if channel.is_closed:
-                logger.warning(f"Channel for queue {queue_name} is closed")
-                # Attempt to recreate the channel if needed
-                if self.connection and not self.connection.is_closed:
-                    logger.info(f"Creating new channel for {queue_name}")
-                    self.channels[queue_name] = await self.connection.channel()
-                    await self.channels[queue_name].set_qos(
-                        prefetch_count=self.config.prefetch_count
-                    )
-                    return self.channels[queue_name]
-                return None
+                logger.warning("Channel for queue %s is closed", queue_name)
+                # Remove the closed channel
+                del self.channels[queue_name]
+
+                # Attempt to recreate the channel if connection is healthy
+                if (
+                    self.connection
+                    and not self.connection.is_closed
+                    and self.connection_healthy
+                ):
+                    try:
+                        logger.debug("Creating new channel for %s", queue_name)
+                        self.channels[queue_name] = await self.connection.channel()
+                        await self.channels[queue_name].set_qos(
+                            prefetch_count=self.config.prefetch_count
+                        )
+                        return self.channels[queue_name]
+                    except Exception as e:
+                        logger.error(
+                            "Failed to recreate channel for %s: %s", queue_name, e
+                        )
+                        # Trigger shutdown if channel creation fails
+                        self._trigger_shutdown()
+                        return None
+                else:
+                    # Connection is not healthy, trigger shutdown
+                    self._trigger_shutdown()
+                    return None
             return channel
         except Exception as e:
-            logger.error(f"Error accessing channel for queue {queue_name}: {e}")
+            logger.error("Error accessing channel for queue %s: %s", queue_name, e)
+            # Trigger shutdown on any channel access error
+            self._trigger_shutdown()
             return None
 
     async def _establish_channel(self, queue_name: str) -> aio_pika.abc.AbstractChannel:
@@ -459,14 +580,14 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         """
         if self.connection is None or self.connection.is_closed:
             logger.warning(
-                f"Cannot create channel for {queue_name}: connection is not available"
+                "Cannot create channel for %s: connection is not available", queue_name
             )
             raise RuntimeError("Connection is not available")
 
-        logger.debug(f"Creating channel for queue {queue_name}")
+        logger.debug("Creating channel for queue %s", queue_name)
         channel = await self.connection.channel()
         await channel.set_qos(prefetch_count=self.config.prefetch_count)
-        logger.debug(f"Created channel for queue {queue_name}")
+        logger.debug("Created channel for queue %s", queue_name)
         return channel
 
     @asynccontextmanager
@@ -481,17 +602,17 @@ 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=(
-                    aio_pika.exceptions.AMQPConnectionError,
-                    aio_pika.exceptions.AMQPChannelError,
+                    AMQPConnectionError,
+                    AMQPChannelError,
                     ConnectionError,
                 ),
             )
 
             # Save in the channels dict for tracking
             self.channels[queue_name] = channel
-            logger.debug(f"Created new channel for queue {queue_name}")
+            logger.debug("Created new channel for queue %s", queue_name)
 
             try:
                 yield channel
@@ -501,7 +622,7 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                 pass
         except aio_pika.exceptions.AMQPError as e:
             logger.error(
-                f"Error creating channel for queue {queue_name} after retries: {e}"
+                "Error creating channel for queue %s after retries: %s", queue_name, e
             )
             raise
 
@@ -510,12 +631,15 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         Creates a new RabbitMQ connection with retry logic.
         """
         try:
-            logger.info("Establishing connection to RabbitMQ")
-            connection = await aio_pika.connect(self.config.url)
-            logger.info("Connected to RabbitMQ successfully")
+            logger.debug("Establishing connection to RabbitMQ")
+            connection = await aio_pika.connect(
+                self.config.url,
+                heartbeat=self.config.connection_heartbeat_interval,
+            )
+            logger.debug("Connected to RabbitMQ successfully")
             return connection
         except Exception as e:
-            logger.error(f"Failed to connect to RabbitMQ: {e}")
+            logger.error("Failed to connect to RabbitMQ: %s", e)
             raise
 
     @asynccontextmanager
@@ -536,9 +660,9 @@ 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=(
-                    aio_pika.exceptions.AMQPConnectionError,
+                    AMQPConnectionError,
                     ConnectionError,
                     OSError,
                     TimeoutError,
@@ -552,14 +676,15 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
                 pass
         except Exception as e:
             logger.error(
-                f"Failed to establish connection to RabbitMQ after retries: {e}"
+                "Failed to establish connection to RabbitMQ after retries: %s", e
             )
             if self.connection:
                 try:
                     await self.connection.close()
                 except Exception as close_error:
                     logger.error(
-                        f"Error closing connection after connect failure: {close_error}"
+                        "Error closing connection after connect failure: %s",
+                        close_error,
                     )
                 self.connection = None
             raise
@@ -572,22 +697,173 @@ class AioPikaMicroserviceConsumer(MessageBusConsumer):
         Get a channel for a specific queue as a context manager.
         This is safer than using get_channel directly as it ensures proper error handling.
         """
-        channel = await self.get_channel(queue_name)
-        if channel is None:
-            if self.connection and not self.connection.is_closed:
-                # Try to create a new channel
-                async with self.create_channel(queue_name) as new_channel:
-                    yield new_channel
-            else:
-                raise RuntimeError(
-                    f"Cannot get channel for queue {queue_name}: no connection available"
+        max_retries = 3
+        retry_delay = 1.0
+
+        for attempt in range(max_retries):
+            try:
+                channel = await self.get_channel(queue_name)
+                if channel is not None:
+                    try:
+                        yield channel
+                        return
+                    finally:
+                        # We don't close the channel here as it's managed by the consumer
+                        pass
+
+                # No channel available, check connection state
+                if (
+                    self.connection
+                    and not self.connection.is_closed
+                    and self.connection_healthy
+                ):
+                    # Try to create a new channel
+                    async with self.create_channel(queue_name) as new_channel:
+                        yield new_channel
+                        return
+                else:
+                    # Connection is not healthy, trigger shutdown
+                    logger.error(
+                        "Connection not healthy while getting channel for %s, triggering shutdown",
+                        queue_name,
+                    )
+                    self._trigger_shutdown()
+                    raise RuntimeError(
+                        f"Cannot get channel for queue {queue_name}: connection is not healthy"
+                    )
+
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    logger.warning(
+                        "Error getting channel for %s, retrying: %s", queue_name, e
+                    )
+                    await self._wait_delay_or_shutdown(retry_delay)
+                    retry_delay *= 2
+                else:
+                    logger.error(
+                        "Failed to get channel for %s after %s attempts: %s",
+                        queue_name,
+                        max_retries,
+                        e,
+                    )
+                    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.
+        This runs as a background task.
+        """
+        while not self.shutdown_event.is_set():
+            try:
+                await self._wait_delay_or_shutdown(
+                    self.config.connection_health_check_interval
                 )
-        else:
+
+                if self.shutdown_event.is_set():
+                    break
+
+                # Check connection health
+                if not await self._is_connection_healthy():
+                    logger.error(
+                        "Connection health check failed, initiating worker shutdown"
+                    )
+                    self.shutdown()
+                    break
+
+            except asyncio.CancelledError:
+                logger.debug("Connection health monitoring cancelled")
+                break
+            except Exception as e:
+                logger.error("Error in connection health monitoring: %s", e)
+                await self._wait_delay_or_shutdown(5)  # Wait before retrying
+
+    async def _is_connection_healthy(self) -> bool:
+        """
+        Check if the connection is healthy.
+        """
+        try:
+            if self.connection is None or self.connection.is_closed:
+                return False
+
+            # Try to create a temporary channel to test connection
+            async with self.connection.channel() as test_channel:
+                # If we can create a channel, connection is healthy
+                return True
+
+        except Exception as e:
+            logger.debug("Connection health check failed: %s", e)
+            return False
+
+    def _trigger_shutdown(self) -> None:
+        """
+        Trigger worker shutdown due to connection loss.
+        """
+        if not self.shutdown_event.is_set():
+            logger.error("Connection lost, initiating worker shutdown")
+            self.connection_healthy = False
+            self.shutdown()
+
+    async def _cleanup_connection(self) -> None:
+        """
+        Clean up existing connection and channels.
+        """
+        # Cancel existing consumers
+        for queue_name, channel in self.channels.items():
             try:
-                yield channel
-            finally:
-                # We don't close the channel here as it's managed by the consumer
-                pass
+                if not channel.is_closed:
+                    # Cancel consumer if we have its tag
+                    if queue_name in self.consumer_tags:
+                        try:
+                            queue = await channel.get_queue(queue_name, ensure=False)
+                            if queue:
+                                await queue.cancel(self.consumer_tags[queue_name])
+                        except Exception as cancel_error:
+                            logger.warning(
+                                "Error cancelling consumer for %s: %s",
+                                queue_name,
+                                cancel_error,
+                            )
+                        del self.consumer_tags[queue_name]
+            except Exception as e:
+                logger.warning("Error cancelling consumer for %s: %s", queue_name, e)
+
+        # Close channels
+        for queue_name, channel in self.channels.items():
+            try:
+                if not channel.is_closed:
+                    await channel.close()
+            except Exception as e:
+                logger.warning("Error closing channel for %s: %s", queue_name, e)
+
+        self.channels.clear()
+
+        # Close connection
+        if self.connection and not self.connection.is_closed:
+            try:
+                await self.connection.close()
+            except Exception as e:
+                logger.warning("Error closing connection: %s", e)
+
+        self.connection = None
+        self.connection_healthy = False
 
 
 def create_message_bus(
@@ -627,11 +903,15 @@ 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
+        connection_heartbeat_interval = 30.0
+        connection_health_check_interval = 10.0
+
         # Connection retry config parameters
         if (
             "connection_retry_max" in query_params
@@ -698,12 +978,31 @@ def create_message_bus(
             except ValueError:
                 pass
 
+        # Heartbeat and health check intervals
+        if "connection_heartbeat_interval" in query_params:
+            try:
+                connection_heartbeat_interval = float(
+                    query_params["connection_heartbeat_interval"][0]
+                )
+            except ValueError:
+                pass
+
+        if "connection_health_check_interval" in query_params:
+            try:
+                connection_health_check_interval = float(
+                    query_params["connection_health_check_interval"][0]
+                )
+            except ValueError:
+                pass
+
         config = AioPikaWorkerConfig(
             url=broker_url,
             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,
         )
 
         return AioPikaMicroserviceConsumer(
@@ -737,37 +1036,59 @@ class ScheduledMessageHandlerCallback:
     ) -> None:
 
         if self.consumer.shutdown_event.is_set():
-            logger.info(
-                f"Shutdown in progress. Requeuing scheduled message for {self.queue_name}"
+            logger.debug(
+                "Shutdown in progress. Requeuing scheduled message for %s",
+                self.queue_name,
             )
             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(
-                    f"Could not requeue scheduled message during shutdown - channel not available"
+                    "Could not requeue scheduled message during shutdown - channel not available"
                 )
             except Exception as e:
                 logger.error(
-                    f"Failed to requeue scheduled message during shutdown: {e}"
+                    "Failed to requeue scheduled message during shutdown: %s", e
                 )
             return
 
+        # Check if connection is healthy before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing scheduled message for %s",
+                self.queue_name,
+            )
+            try:
+                if not self.consumer.connection_healthy:
+                    # Still not healthy, requeue the message
+
+                    await aio_pika_message.reject(requeue=True)
+                    return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue scheduled message due to connection issues: %s",
+                    e,
+                )
+                return
+
         async with self.consumer.lock:
-            task = asyncio.create_task(self.handle_message(aio_pika_message))
+            task = asyncio.create_task(
+                self.handle_message(aio_pika_message),
+                name=f"ScheduledAction-{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(f"Scheduled task for {self.queue_name} was cancelled")
+            logger.warning("Scheduled task for %s was cancelled", self.queue_name)
             return
 
         if (error := task.exception()) is not None:
             logger.exception(
-                f"Error processing scheduled action {self.queue_name}", exc_info=error
+                "Error processing scheduled action %s", self.queue_name, exc_info=error
             )
 
     async def handle_message(
@@ -775,18 +1096,37 @@ class ScheduledMessageHandlerCallback:
     ) -> None:
 
         if self.consumer.shutdown_event.is_set():
-            logger.info(f"Shutdown event set. Requeuing message for {self.queue_name}")
+            logger.debug(
+                "Shutdown event set. Requeuing message for %s", self.queue_name
+            )
             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(
-                    f"Could not requeue message during shutdown - channel not available"
+                    "Could not requeue message during shutdown - channel not available"
                 )
             except Exception as e:
-                logger.error(f"Failed to requeue message during shutdown: {e}")
+                logger.error("Failed to requeue message during shutdown: %s", e)
+                return
+
+        # Check connection health before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing scheduled message for %s",
+                self.queue_name,
+            )
+            try:
+
+                await aio_pika_message.reject(requeue=True)
+                return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue scheduled message due to connection issues: %s",
+                    e,
+                )
                 return
 
         sig = inspect.signature(self.scheduled_action.callable)
@@ -797,7 +1137,8 @@ class ScheduledMessageHandlerCallback:
                     self.scheduled_action,
                     (ScheduleDispatchData(int(aio_pika_message.body.decode("utf-8"))),),
                     {},
-                )
+                ),
+                name=f"ScheduledAction-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
             )
 
         elif len(sig.parameters) == 0:
@@ -806,7 +1147,8 @@ class ScheduledMessageHandlerCallback:
                     self.scheduled_action,
                     (),
                     {},
-                )
+                ),
+                name=f"ScheduledAction-{self.queue_name}-handle-message-{aio_pika_message.message_id}",
             )
         else:
             logger.warning(
@@ -823,7 +1165,7 @@ class ScheduledMessageHandlerCallback:
         except Exception as e:
 
             logger.exception(
-                f"Error processing scheduled action {self.queue_name}: {e}"
+                "Error processing scheduled action %s: %s", self.queue_name, e
             )
 
     async def run_with_context(
@@ -832,18 +1174,22 @@ class ScheduledMessageHandlerCallback:
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
     ) -> None:
-        async with self.consumer.uow_context_provider(
-            AppTransactionContext(
-                controller_member_reflect=scheduled_action.controller_member,
-                transaction_data=SchedulerTransactionData(
-                    scheduled_to=datetime.now(UTC),
-                    cron_expression=scheduled_action.spec.cron,
-                    triggered_at=datetime.now(UTC),
-                ),
-            )
-        ):
 
-            await scheduled_action.callable(*args, **kwargs)
+        with provide_shutdown_state(self.consumer.shutdown_state):
+            async with self.consumer.uow_context_provider(
+                AppTransactionContext(
+                    controller_member_reflect=scheduled_action.controller_member,
+                    transaction_data=SchedulerTransactionData(
+                        task_name=scheduled_action.spec.name
+                        or scheduled_action.callable.__qualname__,
+                        scheduled_to=datetime.now(UTC),
+                        cron_expression=scheduled_action.spec.cron,
+                        triggered_at=datetime.now(UTC),
+                    ),
+                )
+            ):
+
+                await scheduled_action.callable(*args, **kwargs)
 
 
 class MessageHandlerCallback:
@@ -865,36 +1211,60 @@ class MessageHandlerCallback:
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
     ) -> None:
         if self.consumer.shutdown_event.is_set():
-            logger.info(
-                f"Shutdown in progress. Requeuing message for {self.queue_name}"
+            logger.debug(
+                "Shutdown in progress. Requeuing message for %s", self.queue_name
             )
             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(
-                    f"Could not requeue message during shutdown - channel not available"
+                    "Could not requeue message during shutdown - channel not available"
                 )
             except Exception as e:
-                logger.error(f"Failed to requeue message during shutdown: {e}")
+                logger.error("Failed to requeue message during shutdown: %s", e)
             return
 
+        # Check if connection is healthy before processing
+        if not self.consumer.connection_healthy:
+            logger.warning(
+                "Connection not healthy, requeuing message for %s", self.queue_name
+            )
+            try:
+                if not self.consumer.connection_healthy:
+                    # Still not healthy, requeue the message
+
+                    await aio_pika_message.reject(requeue=True)
+                    return
+            except Exception as e:
+                logger.error(
+                    "Failed to requeue message due to connection issues: %s", e
+                )
+                return
+
         async with self.consumer.lock:
-            task = asyncio.create_task(self.handle_message(aio_pika_message))
+            task = asyncio.create_task(
+                self.handle_message(aio_pika_message),
+                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(f"Task for queue {self.queue_name} was cancelled")
-            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(
-                f"Error processing message for queue {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
@@ -904,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,
@@ -919,33 +1291,35 @@ class MessageHandlerCallback:
         """
         message_id = aio_pika_message.message_id or str(uuid.uuid4())
 
-        # If auto_ack is enabled, we cannot retry the message through RabbitMQ reject mechanism
-        if self.message_handler.spec.auto_ack:
-            if requeue:
-                logger.warning(
-                    f"Message {message_id} ({self.queue_name}) cannot be requeued because auto_ack is enabled"
-                )
-            return
-
         try:
             # Check if we should retry with backoff
             if (
                 not requeue
-                and self.message_handler.spec.requeue_on_exception
+                and self.message_handler.spec.nack_on_exception
                 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:
                     logger.warning(
-                        f"Message {message_id} ({self.queue_name}) failed after {retry_count} retries, "
-                        f"dead-lettering: {str(exception)}"
+                        "Message %s (%s) failed after %s retries, dead-lettering: %s",
+                        message_id,
+                        self.queue_name,
+                        retry_count,
+                        str(exception),
                     )
                     # Dead-letter the message after max retries
-                    async with self.consumer.get_channel_ctx(self.queue_name):
+                    try:
                         await aio_pika_message.reject(requeue=False)
+                    except Exception as e:
+                        logger.error(
+                            "Failed to dead-letter message %s: %s", message_id, e
+                        )
                     return
 
                 # Calculate delay for this retry attempt
@@ -961,9 +1335,14 @@ class MessageHandlerCallback:
 
                 delay = min(delay, retry_config.max_delay)
 
-                logger.info(
-                    f"Message {message_id} ({self.queue_name}) failed with {str(exception)}, "
-                    f"retry {retry_count+1}/{retry_config.max_retries} scheduled in {delay:.2f}s"
+                logger.warning(
+                    "Message %s (%s) failed with %s, retry %s/%s scheduled in %.2fs",
+                    message_id,
+                    self.queue_name,
+                    str(exception),
+                    retry_count + 1,
+                    retry_config.max_retries,
+                    delay,
                 )
 
                 # Store retry state for this message
@@ -974,36 +1353,50 @@ 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
-                async with self.consumer.get_channel_ctx(self.queue_name):
+                try:
                     await aio_pika_message.ack()
+                except Exception as e:
+                    logger.error(
+                        "Failed to acknowledge message %s for retry: %s", message_id, e
+                    )
                 return
 
             # Standard reject without retry or with immediate requeue
-            async with self.consumer.get_channel_ctx(self.queue_name):
+            try:
+                await self._wait_delay_or_shutdown(
+                    requeue_timeout
+                )  # Optional delay before requeueing
                 await aio_pika_message.reject(requeue=requeue)
                 if requeue:
-                    logger.info(
-                        f"Message {message_id} ({self.queue_name}) requeued for immediate retry"
+                    logger.warning(
+                        "Message %s (%s) requeued for immediate retry",
+                        message_id,
+                        self.queue_name,
                     )
                 else:
-                    logger.info(
-                        f"Message {message_id} ({self.queue_name}) rejected without requeue"
+                    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)
 
-        except RuntimeError as e:
-            logger.error(
-                f"Error rejecting message {message_id} ({self.queue_name}): {e}"
-            )
         except Exception as e:
             logger.exception(
-                f"Unexpected error rejecting message {message_id} ({self.queue_name}): {e}"
+                "Unexpected error in handle_reject_message for %s (%s): %s",
+                message_id,
+                self.queue_name,
+                e,
             )
 
     async def _delayed_retry(
@@ -1018,7 +1411,7 @@ class MessageHandlerCallback:
 
         Args:
             aio_pika_message: The original message
-            delay: Delay in seconds before retry
+            delay: Delay in seconds before retrying
             retry_count: The current retry count (after increment)
             exception: The exception that caused the failure
         """
@@ -1026,7 +1419,7 @@ class MessageHandlerCallback:
 
         try:
             # Wait for the backoff delay
-            await asyncio.sleep(delay)
+            await self._wait_delay_or_shutdown(delay)
 
             # Get message body and properties for republishing
             message_body = aio_pika_message.body
@@ -1043,32 +1436,64 @@ class MessageHandlerCallback:
             if message_id in self.retry_state:
                 del self.retry_state[message_id]
 
-            # Republish the message to the same queue
-            async with self.consumer.get_channel_ctx(self.queue_name) as channel:
-                exchange = await RabbitmqUtils.get_main_exchange(
-                    channel=channel,
-                    exchange_name=self.consumer.config.exchange,
-                )
+            # Republish the message to the same queue with retry logic
+            max_attempts = 3
+            for attempt in range(max_attempts):
+                try:
+                    async with self.consumer.get_channel_ctx(
+                        self.queue_name
+                    ) as channel:
+                        exchange = await RabbitmqUtils.get_main_exchange(
+                            channel=channel,
+                            exchange_name=self.consumer.config.exchange,
+                        )
 
-                await exchange.publish(
-                    aio_pika.Message(
-                        body=message_body,
-                        headers=headers,
-                        message_id=message_id,
-                        content_type=aio_pika_message.content_type,
-                        content_encoding=aio_pika_message.content_encoding,
-                        delivery_mode=aio_pika_message.delivery_mode,
-                    ),
-                    routing_key=self.routing_key,
-                )
+                        await exchange.publish(
+                            aio_pika.Message(
+                                body=message_body,
+                                headers=headers,
+                                message_id=message_id,
+                                content_type=aio_pika_message.content_type,
+                                content_encoding=aio_pika_message.content_encoding,
+                                delivery_mode=aio_pika_message.delivery_mode,
+                            ),
+                            routing_key=self.routing_key,
+                        )
 
-                logger.info(
-                    f"Message {message_id} ({self.queue_name}) republished for retry {retry_count}"
-                )
+                        logger.warning(
+                            "Message %s (%s) republished for retry %s",
+                            message_id,
+                            self.queue_name,
+                            retry_count,
+                        )
+                        return
+
+                except Exception as e:
+                    if attempt < max_attempts - 1:
+                        logger.warning(
+                            "Failed to republish message %s (attempt %s): %s",
+                            message_id,
+                            attempt + 1,
+                            e,
+                        )
+                        await self._wait_delay_or_shutdown(
+                            (1.0 * (attempt + 1))
+                        )  # Exponential backoff
+                    else:
+                        logger.error(
+                            "Failed to republish message %s after %s attempts: %s",
+                            message_id,
+                            max_attempts,
+                            e,
+                        )
+                        raise
 
         except Exception as e:
             logger.exception(
-                f"Failed to execute delayed retry for message {message_id} ({self.queue_name}): {e}"
+                "Failed to execute delayed retry for message %s (%s): %s",
+                message_id,
+                self.queue_name,
+                e,
             )
             # If we fail to republish, try to dead-letter the original message
             try:
@@ -1077,6 +1502,9 @@ class MessageHandlerCallback:
             except Exception:
                 pass
 
+    async def _wait_delay_or_shutdown(self, delay: float) -> None:
+        await self.consumer._wait_delay_or_shutdown(delay)
+
     async def handle_message(
         self, aio_pika_message: aio_pika.abc.AbstractIncomingMessage
     ) -> None:
@@ -1091,124 +1519,224 @@ 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)
 
-        incoming_message_spec = MessageHandler.get_message_incoming(handler)
-        assert incoming_message_spec is not None
+        built_message = AioPikaMessage(aio_pika_message, message_type)
 
-        async with self.consumer.uow_context_provider(
-            AppTransactionContext(
-                controller_member_reflect=handler_data.controller_member,
-                transaction_data=MessageBusTransactionData(
-                    message=builded_message,
-                    topic=routing_key,
-                ),
-            )
+        incoming_message_spec = MessageHandler.get_last(handler)
+        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
         ):
-            ctx: AsyncContextManager[Any]
-            if incoming_message_spec.timeout is not None:
-                ctx = asyncio.timeout(incoming_message_spec.timeout)
-            else:
-                ctx = none_context()
-            async with ctx:
-                try:
-                    with provide_bus_message_controller(
-                        AioPikaMessageBusController(aio_pika_message)
-                    ):
-                        await handler(builded_message)
-                    if not incoming_message_spec.auto_ack:
-                        with suppress(aio_pika.MessageProcessError):
-                            # Use channel context for acknowledgement
-                            async with self.consumer.get_channel_ctx(self.queue_name):
-                                await aio_pika_message.ack()
-                except BaseException as base_exc:
-                    # Get message id for logging
-                    message_id = aio_pika_message.message_id or str(uuid.uuid4())
-
-                    # 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(
-                                f"Error processing exception handler for message {message_id}: {base_exc} | {nested_exc}"
+            async with self.consumer.uow_context_provider(
+                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=built_message,
+                        topic=routing_key,
+                    ),
+                )
+            ):
+                maybe_timeout_ctx: AsyncContextManager[Any]
+                if incoming_message_spec.timeout is not None:
+                    maybe_timeout_ctx = asyncio.timeout(incoming_message_spec.timeout)
+                else:
+                    maybe_timeout_ctx = none_context()
+
+                start_time = time.perf_counter()
+                async with maybe_timeout_ctx:
+                    try:
+                        with provide_bus_message_controller(
+                            AioPikaMessageBusController(aio_pika_message)
+                        ):
+                            try:
+                                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,
+                                    {
+                                        "message_id": aio_pika_message.message_id
+                                        or "unknown",
+                                        "routing_key": routing_key,
+                                    },
+                                )
+                                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_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()
+                            )
+                            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 ""
                             )
-                    else:
-                        logger.exception(
-                            f"Error processing message {message_id} on topic {routing_key}: {str(base_exc)}"
-                        )
 
-                    # Handle rejection with retry logic
-                    if incoming_message_spec.requeue_on_exception:
-                        # Use our retry with backoff mechanism
-                        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,
-                        )
-                    else:
-                        # Message shouldn't be retried, reject it
-                        await self.handle_reject_message(
-                            aio_pika_message, requeue=False, exception=base_exc
-                        )
-                else:
-                    # 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 {}
-                    if "x-retry-count" in headers:
-                        retry_count = int(str(headers.get("x-retry-count", 0)))
-                        logger.info(
-                            f"Message {message_id}#{self.queue_name} processed successfully after {retry_count} retries"
-                        )
-                    else:
-                        logger.info(
-                            f"Message {message_id}#{self.queue_name} processed successfully"
+                            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}"
                         )
 
 
@@ -1248,68 +1776,73 @@ class MessageBusWorker:
     async def start_async(self) -> None:
         all_message_handlers_set: MESSAGE_HANDLER_DATA_SET = set()
         all_scheduled_actions_set: SCHEDULED_ACTION_DATA_SET = set()
-        async with self.lifecycle():
-            for instance_class in self.app.controllers:
-                controller = MessageBusController.get_messagebus(instance_class)
-
-                if controller is None:
-                    continue
-
-                instance: Any = self.container.get_by_type(instance_class)
-
-                factory = controller.get_messagebus_factory()
-                handlers, schedulers = factory(instance)
-
-                message_handler_data_map: dict[str, MessageHandlerData] = {}
-                all_scheduled_actions_set.update(schedulers)
-                for handler_data in handlers:
-                    message_type = handler_data.spec.message_type
-                    topic = message_type.MESSAGE_TOPIC
-
-                    # Filter handlers by name if specified
-                    if (
-                        self.handler_names is not None
-                        and handler_data.spec.name is not None
-                    ):
-                        if handler_data.spec.name not in self.handler_names:
-                            continue
-                    elif (
-                        self.handler_names is not None
-                        and handler_data.spec.name is None
-                    ):
-                        # Skip handlers without names when filtering is requested
-                        continue
+        with providing_app_type("worker"):
+            async with self.lifecycle():
+                for instance_class in self.app.controllers:
+                    controller = MessageBusController.get_last(instance_class)
 
-                    if (
-                        topic in message_handler_data_map
-                        and message_type.MESSAGE_TYPE == "task"
-                    ):
-                        logger.warning(
-                            "Task handler for topic '%s' already registered. Skipping"
-                            % topic
-                        )
+                    if controller is None:
                         continue
-                    message_handler_data_map[topic] = handler_data
-                    all_message_handlers_set.add(handler_data)
 
-            broker_backend = get_message_broker_backend_from_url(url=self.backend_url)
+                    instance: Any = self.container.get_by_type(instance_class)
+
+                    factory = controller.get_messagebus_factory()
+                    handlers, schedulers = factory(instance)
+
+                    message_handler_data_map: dict[str, MessageHandlerData] = {}
+                    all_scheduled_actions_set.update(schedulers)
+                    for handler_data in handlers:
+                        message_type = handler_data.spec.message_type
+                        topic = message_type.MESSAGE_TOPIC
+
+                        # Filter handlers by name if specified
+                        if (
+                            self.handler_names is not None
+                            and handler_data.spec.name is not None
+                        ):
+                            if handler_data.spec.name not in self.handler_names:
+                                continue
+                        elif (
+                            self.handler_names is not None
+                            and handler_data.spec.name is None
+                        ):
+                            # Skip handlers without names when filtering is requested
+                            continue
+
+                        if (
+                            topic in message_handler_data_map
+                            and message_type.MESSAGE_TYPE == "task"
+                        ):
+                            logger.warning(
+                                "Task handler for topic '%s' already registered. Skipping"
+                                % topic
+                            )
+                            continue
+                        message_handler_data_map[topic] = handler_data
+                        all_message_handlers_set.add(handler_data)
 
-            consumer = self._consumer = create_message_bus(
-                broker_url=self.broker_url,
-                broker_backend=broker_backend,
-                scheduled_actions=all_scheduled_actions_set,
-                message_handler_set=all_message_handlers_set,
-                uow_context_provider=self.uow_context_provider,
-            )
+                broker_backend = get_message_broker_backend_from_url(
+                    url=self.backend_url
+                )
 
-            await consumer.consume()
+                consumer = self._consumer = create_message_bus(
+                    broker_url=self.broker_url,
+                    broker_backend=broker_backend,
+                    scheduled_actions=all_scheduled_actions_set,
+                    message_handler_set=all_message_handlers_set,
+                    uow_context_provider=self.uow_context_provider,
+                )
+
+                await consumer.consume()
 
     def start_sync(self) -> None:
 
         def on_shutdown(loop: asyncio.AbstractEventLoop) -> None:
-            logger.info("Shutting down - signal received")
+            logger.warning("Shutting down - signal received")
             # Schedule the shutdown to run in the event loop
-            asyncio.create_task(self._graceful_shutdown())
+            asyncio.create_task(
+                self._graceful_shutdown(), name="Worker-Graceful-Shutdown"
+            )
             # wait until the shutdown is complete
 
         with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
@@ -1317,15 +1850,22 @@ class MessageBusWorker:
             loop.add_signal_handler(signal.SIGINT, on_shutdown, loop)
             # Add graceful shutdown handler for SIGTERM as well
             loop.add_signal_handler(signal.SIGTERM, on_shutdown, loop)
-            runner.run(self.start_async())
+            try:
+                runner.run(self.start_async())
+            except Exception as e:
+                logger.critical("Worker failed to start due to connection error: %s", e)
+                # Exit with error code 1 to indicate startup failure
+                import sys
+
+                sys.exit(1)
 
     async def _graceful_shutdown(self) -> None:
         """Handles graceful shutdown process"""
-        logger.info("Initiating graceful shutdown sequence")
+        logger.warning("Initiating graceful shutdown sequence")
         # Use the comprehensive close method that handles shutdown, task waiting and connection cleanup
 
         self.consumer.shutdown()
-        logger.info("Graceful shutdown completed")
+        logger.warning("Graceful shutdown completed")
 
 
 class AioPikaMessageBusController(BusMessageController):
@@ -1375,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:
         """
@@ -1404,13 +1946,14 @@ class AioPikaMessageBusController(BusMessageController):
                     float(delay),
                     retry_count + 1,
                     None,  # No specific exception
-                )
+                ),
+                name=f"MessageHandler-{callback.queue_name}-delayed-retry-{self.aio_pika_message.message_id or 'unknown'}-{int(time.time())}",
             )
 
             # Acknowledge the current message since we'll republish
             await self.aio_pika_message.ack()
 
         except Exception as e:
-            logger.exception(f"Failed to schedule retry_later: {e}")
+            logger.exception("Failed to schedule retry_later: %s", e)
             # Fall back to immediate retry
             await self.aio_pika_message.reject(requeue=True)