fix(queue): use queue as ctx, always close connection

Author: peshence

polytope_server/broker/broker.py

@@ -30,8 +30,8 @@
 class Broker:
     def __init__(self, config: dict):
 
-        queue_config = config.get("queue")
-        self.queue = queue.create_queue(queue_config)
+        self.queue_config = config.get("queue")
+        self.queue = None
 
         self.max_queue_size = config.get("deployment", {}).get("worker", {}).get("replicas", 40)
 
@@ -47,8 +47,14 @@ def run(self):
         logging.info("Starting broker...")
         logging.info("Maximum Queue Size: {}".format(self.max_queue_size))
 
-        while not time.sleep(self.scheduling_interval):
-            self.check_requests()
+        with queue.create_queue(self.queue_config) as q:
+            self.queue = q
+            try:
+                while not time.sleep(self.scheduling_interval):
+                    self.check_requests()
+            except KeyboardInterrupt:
+                logging.info("Broker shutdown requested")
+                raise
 
     def check_requests(self):
 

polytope_server/common/queue/queue.py

@@ -71,6 +71,15 @@ def count(self) -> int:
     def get_type(self) -> str:
         """Get the implementation type"""
 
+    def __enter__(self):
+        """Enter context manager"""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and close connection"""
+        self.close_connection()
+        return False
+
 
 queue_types = {"rabbitmq": RabbitmqQueue, "sqs": SQSQueue}
 

polytope_server/worker/worker.py

@@ -205,12 +205,12 @@ def handle_termination(group: aio.TaskGroup) -> None:
             loop.remove_signal_handler(signal.SIGTERM)
 
     def run(self):
-        self.queue = polytope_queue.create_queue(self.config.get("queue"))
+        with polytope_queue.create_queue(self.config.get("queue")) as queue:
+            self.queue = queue
+            self.update_status("idle", time_spent=0)
 
-        self.update_status("idle", time_spent=0)
-
-        with ThreadPoolExecutor(max_workers=1) as executor:
-            aio.run(self.schedule(executor))
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                aio.run(self.schedule(executor))
 
     def process_request(
         self,