Adding a killing feature for jobs (mrq/{job.py,worker.py)

Author: DasFranck

mrq/job.py

@@ -249,6 +249,14 @@ def abort(self):
         self._attach_original_exception(exc)
         raise exc
 
+    def kill(self):
+        """ Kill the current job """
+        context.connections.redis.rpush("{}:wcmd:{}".format(context.get_current_config()["redis_prefix"],
+                                                            context.get_current_worker()),
+                                        "kill {}".format(self.id))
+        self._save_status("killed")
+        pass
+
     def cancel(self):
         """ Markes the current job as cancelled. Doesn't interrupt it. """
         self._save_status("cancel")
@@ -321,15 +329,15 @@ def perform(self):
             # pylint: disable=protected-access
 
             gevent.sleep(0)
-            current_greenlet = gevent.getcurrent()
+            self.current_greenlet = gevent.getcurrent()
             t = (datetime.datetime.utcnow() - self.datestarted).total_seconds()
 
             context.log.debug(
                 "Job %s success: %0.6fs total, %0.6fs in greenlet, %s switches" %
                 (self.id,
                  t,
-                 current_greenlet._trace_time,
-                 current_greenlet._trace_switches - 1)
+                 self.current_greenlet._trace_time,
+                 self.current_greenlet._trace_switches - 1)
             )
 
         else:
@@ -454,13 +462,13 @@ def _save_status(self, status, updates=None, exception=False, w=None, j=None):
             db_updates["totaltime"] = (now - self.datestarted).total_seconds()
 
         if context.get_current_config().get("trace_greenlets"):
-            current_greenlet = gevent.getcurrent()
+            self.current_greenlet = gevent.getcurrent()
 
             # TODO are we sure the current job is doing the save_status() on itself?
-            if hasattr(current_greenlet, "_trace_time"):
+            if hasattr(self.current_greenlet, "_trace_time"):
                 # pylint: disable=protected-access
-                db_updates["time"] = current_greenlet._trace_time
-                db_updates["switches"] = current_greenlet._trace_switches
+                db_updates["time"] = self.current_greenlet._trace_time
+                db_updates["switches"] = self.current_greenlet._trace_switches
 
         if exception:
             trace = traceback.format_exc()

mrq/worker.py

@@ -328,6 +328,18 @@ def report_worker(self, w=0):
         except Exception as e:  # pylint: disable=broad-except
             self.log.debug("Worker report failed: %s" % e)
 
+    def greenlet_command_handler(self):
+        """
+        This greenlet is used to execute commands directly in the worker
+        """
+        while True:
+            command = self.redis.blpop("{}:wcmd:{}".format(get_current_config()["redis_prefix"],
+                                                           self.id)
+                                       ).split(" ")
+            if command[0] == "kill":
+                Job(command[1]).current_greenlet.kill(block=True)
+            pass
+
     def greenlet_admin(self):
         """ This greenlet is used to get status information about the worker
             when --admin_port was given
@@ -436,6 +448,9 @@ def work_init(self):
         if self.config["admin_port"]:
             self.greenlets["admin"] = gevent.spawn(self.greenlet_admin)
 
+        # ehould add a condition using the config
+        self.greenlets["command_handler"] = gevent.spawn(self.greenlet_command_handler)
+
         self.install_signal_handlers()
 
     def work_loop(self, max_jobs=None, max_time=None):