the issue is dameon rep socket is now blocking, so we wait

Author: sudoRicheek

src/cortex/discovery/daemon.py

@@ -10,7 +10,6 @@
 
 import contextlib
 import os
-import threading
 
 import zmq
 
@@ -57,9 +56,8 @@ def __init__(
         self._context: zmq.Context | None = None
         self._socket: zmq.Socket | None = None
 
-        # Control flags
+        # Control flag
         self._running = False
-        self._shutdown_event = threading.Event()
 
     def _ensure_ipc_path(self) -> None:
         """Ensure the IPC socket directory exists."""
@@ -90,7 +88,6 @@ def start(self) -> None:
         self._socket.setsockopt(zmq.LINGER, 0)  # Immediate shutdown
 
         self._running = True
-        self._shutdown_event.clear()
 
         logger.info("=" * 50)
         logger.info("DISCOVERY DAEMON STARTED")
@@ -102,24 +99,27 @@ def start(self) -> None:
         except KeyboardInterrupt:
             logger.info("Received interrupt signal")
         finally:
-            self.stop()
+            self._cleanup()
 
     def _run_loop(self) -> None:
         """Main event loop."""
-        while self._running and not self._shutdown_event.is_set():
+        while self._running:
             try:
-                # Try to receive a request
+                # Try to receive a request (blocks up to RCVTIMEO)
                 try:
                     request_bytes = self._socket.recv(copy=False)
 
                     # Process the request
                     response = self._handle_request(request_bytes)
                     self._socket.send(response.to_bytes())
                 except zmq.Again:
-                    # No message available, continue
+                    # Timeout, check _running and continue
                     continue
 
             except Exception as e:
+                if not self._running:
+                    # We're shutting down, exit cleanly
+                    break
                 logger.error(f"Error in discovery loop: {e}")
                 # Send error response if we received a request
                 try:
@@ -253,12 +253,8 @@ def _handle_shutdown(self) -> DiscoveryResponse:
         logger.info("SHUTDOWN command received")
         return DiscoveryResponse(status=DiscoveryStatus.OK, message="Shutting down")
 
-    def stop(self) -> None:
-        """Stop the discovery daemon."""
-        logger.info("Stopping discovery daemon")
-        self._shutdown_event.set()
-        self._running = False
-
+    def _cleanup(self) -> None:
+        """Clean up resources."""
         if self._socket:
             try:
                 self._socket.close()
@@ -284,6 +280,11 @@ def stop(self) -> None:
         logger.info("DISCOVERY DAEMON STOPPED")
         logger.info("=" * 50)
 
+    def stop(self) -> None:
+        """Stop the discovery daemon."""
+        logger.info("Stopping discovery daemon")
+        self._running = False
+
 
 def main():
     """Entry point for the discovery daemon."""

tests/conftest.py

@@ -77,8 +77,7 @@ def run_daemon():
     def stop(self) -> None:
         """Stop the discovery daemon."""
         if self._daemon:
-            self._daemon._running = False
-            self._daemon._shutdown_event.set()
+            self._daemon.stop()
 
         # Wait for thread to finish (daemon cleans up its own context)
         if self._thread: