Fix agent shutdown

Author: sylvinus

mrq/agent.py

@@ -31,9 +31,11 @@ def work(self):
 
         self.pool.start()
 
-        self.pool.wait()
-
-        connections.mongodb_jobs.mrq_agents.delete_one({"_id": self.id})
+        try:
+            self.pool.wait()
+        finally:
+            self.shutdown_now()
+            connections.mongodb_jobs.mrq_agents.delete_one({"_id": self.id})
 
     def shutdown_now(self):
         self.pool.terminate()
@@ -81,7 +83,8 @@ def get_agent_report(self):
             "available_cpu": get_current_config()["available_cpu"],
             "available_memory": get_current_config()["available_memory"],
             "worker_group": self.worker_group,
-
+            "datereported": datetime.datetime.utcnow(),
+            "dateexpires": datetime.datetime.utcnow() + datetime.timedelta(seconds=(self.config["report_interval"] * 3) + 5)
         }
         return report
 
@@ -184,8 +187,7 @@ def orchestrate(self):
                 connections.mongodb_jobs.mrq_agents.update_one({"_id": agent["_id"]}, {"$set": {
                     "desired_workers": agent["new_desired_workers"],
                     "free_cpu": agent["free_cpu"],
-                    "free_memory": agent["free_memory"],
-                    "datereported": datetime.datetime.utcnow()
+                    "free_memory": agent["free_memory"]
                 }})
 
     def get_desired_workers_for_group(self, group):
@@ -206,4 +208,3 @@ def fetch_worker_group_agents(self):
 
     def fetch_worker_group_definition(self):
         return connections.mongodb_jobs.mrq_workergroups.find_one({"_id": self.worker_group})
-

mrq/worker.py

@@ -151,7 +151,9 @@ def ensure_indexes(self):
             [("hash", 1)], unique=True, background=False)
 
         self.mongodb_jobs.mrq_agents.ensure_index(
-            [("datereported", 1)], background=False, expireAfterSeconds=300)
+            [("datereported", 1)], background=False)
+        self.mongodb_jobs.mrq_agents.ensure_index(
+            [("dateexpires", 1)], background=False, expireAfterSeconds=0)
         self.mongodb_jobs.mrq_agents.ensure_index(
             [("worker_group", 1)], background=False)
 

tests/conftest.py

@@ -97,7 +97,7 @@ def stop(self, force=False, timeout=None, block=True, sig=15):
 
         if self.process is not None:
 
-            print("Test sending signal %s to %s" % (sig, self.process.pid))
+            print("Sending signal %s to pid %s" % (sig, self.process.pid))
             os.kill(self.process.pid, sig)
 
             # When sending a sigkill to the process, we also want to kill the

tests/test_agent.py

@@ -153,14 +153,14 @@ def test_agent_process(worker):
 
     connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xxx", "profiles": [
         {
-            "command": "mrq-worker a",
+            "command": "mrq-worker a --report_interval=1",
             "memory": 100,
             "cpu": 100,
             "min_count": 1
         }
     ]})
 
-    time.sleep(3)
+    time.sleep(7)
 
     assert connections.mongodb_jobs.mrq_workers.count() == 1
     w = connections.mongodb_jobs.mrq_workers.find_one()
@@ -179,6 +179,8 @@ def test_agent_process(worker):
 
     worker.stop(deps=False)
 
+    time.sleep(2)
+
     assert connections.mongodb_jobs.mrq_agents.count() == 0
 
     worker.stop_deps()