REFACTOR external service and improve websocket training

shared/messaging/rabbitmq.py

@@ -7,6 +7,7 @@ from typing import Dict, Any, Callable, Optional
 from datetime import datetime, date
 import uuid
 import structlog
+from contextlib import suppress
 
 try:
     import aio_pika
@@ -17,6 +18,50 @@ except ImportError:
     
 logger = structlog.get_logger()
 
+class HeartbeatMonitor:
+    """Monitor to ensure heartbeats are processed during heavy operations"""
+    
+    def __init__(self, client):
+        self.client = client
+        self._monitor_task = None
+        self._should_monitor = False
+        
+    async def start_monitoring(self):
+        """Start heartbeat monitoring task"""
+        if self._monitor_task and not self._monitor_task.done():
+            return
+            
+        self._should_monitor = True
+        self._monitor_task = asyncio.create_task(self._monitor_loop())
+        
+    async def stop_monitoring(self):
+        """Stop heartbeat monitoring task"""
+        self._should_monitor = False
+        if self._monitor_task and not self._monitor_task.done():
+            self._monitor_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._monitor_task
+    
+    async def _monitor_loop(self):
+        """Monitor loop that periodically yields control for heartbeat processing"""
+        while self._should_monitor:
+            # Yield control to allow heartbeat processing
+            await asyncio.sleep(0.1)
+            
+            # Verify connection is still alive
+            if self.client.connection and not self.client.connection.is_closed:
+                # Check if connection is still responsive
+                try:
+                    # This is a lightweight check to ensure the connection is responsive
+                    pass  # The heartbeat mechanism in aio_pika handles this internally
+                except Exception as e:
+                    logger.warning("Connection check failed", error=str(e))
+                    self.client.connected = False
+                    break
+            else:
+                logger.warning("Connection is closed, stopping monitor")
+                break
+
 def json_serializer(obj):
     """JSON serializer for objects not serializable by default json code"""
     if isinstance(obj, (datetime, date)):
@@ -42,6 +87,7 @@ class RabbitMQClient:
         self.connected = False
         self._reconnect_attempts = 0
         self._max_reconnect_attempts = 5
+        self.heartbeat_monitor = HeartbeatMonitor(self)
         
     async def connect(self):
         """Connect to RabbitMQ with retry logic"""
@@ -52,14 +98,17 @@ class RabbitMQClient:
         try:
             self.connection = await connect_robust(
                 self.connection_url,
-                heartbeat=30,
-                connection_attempts=3
+                heartbeat=600  # Increase heartbeat to 600 seconds (10 minutes) to prevent timeouts
             )
             self.channel = await self.connection.channel()
             await self.channel.set_qos(prefetch_count=100)  # Performance optimization
             
             self.connected = True
             self._reconnect_attempts = 0
+            
+            # Start heartbeat monitoring
+            await self.heartbeat_monitor.start_monitoring()
+            
             logger.info("Connected to RabbitMQ", service=self.service_name)
             return True
             
@@ -75,11 +124,28 @@ class RabbitMQClient:
             return False
     
     async def disconnect(self):
-        """Disconnect from RabbitMQ"""
-        if self.connection and not self.connection.is_closed:
-            await self.connection.close()
+        """Disconnect from RabbitMQ with proper channel cleanup"""
+        try:
+            # Stop heartbeat monitoring first
+            await self.heartbeat_monitor.stop_monitoring()
+
+            # Close channel before connection to avoid "unexpected close" warnings
+            if self.channel and not self.channel.is_closed:
+                await self.channel.close()
+                logger.debug("RabbitMQ channel closed", service=self.service_name)
+
+            # Then close connection
+            if self.connection and not self.connection.is_closed:
+                await self.connection.close()
+                logger.info("Disconnected from RabbitMQ", service=self.service_name)
+
+            self.connected = False
+
+        except Exception as e:
+            logger.warning("Error during RabbitMQ disconnect",
+                          service=self.service_name,
+                          error=str(e))
             self.connected = False
-            logger.info("Disconnected from RabbitMQ", service=self.service_name)
     
     async def ensure_connected(self) -> bool:
         """Ensure connection is active, reconnect if needed"""