Better management of known queues

Author: sylvinus

mrq/basetasks/cleaning.py

@@ -205,13 +205,33 @@ class CleanKnownQueues(Task):
     def run(self, params):
 
         max_age = int(params.get("max_age") or (7 * 86400))
+        pretend = bool(params.get("pretend"))
+        check_mongo = bool(params.get("check_mongo"))
 
         known_queues = Queue.redis_known_queues()
 
+        removed_queues = []
+
+        queues_from_config = Queue.all_known_from_config()
+
+        print "Found %s known queues & %s from config" % (len(known_queues), len(queues_from_config))
+
         # Only clean queues older than N days
         time_threshold = time.time() - max_age
         for queue, time_last_used in known_queues.iteritems():
+            if queue in queues_from_config:
+                continue
             if time_last_used < time_threshold:
                 q = Queue(queue, add_to_known_queues=False)
-                if q.size() == 0:
-                    q.remove_from_known_queues()
+                size = q.size()
+                if check_mongo:
+                    size += connections.mongodb_jobs.mrq_jobs.count({"queue": queue})
+                if size == 0:
+                    removed_queues.append(queue)
+                    print "Removing empty queue '%s' from known queues ..." % queue
+                    if not pretend:
+                        q.remove_from_known_queues()
+
+        print "Cleaned %s queues" % len(removed_queues)
+
+        return removed_queues

mrq/basetasks/utils.py

@@ -132,7 +132,7 @@ def perform_action(self, action, query, destination_queue):
 
                     # Between these two lines, jobs can become "lost" too.
 
-                    Queue(destination_queue or queue).enqueue_job_ids(
+                    Queue(destination_queue or queue, add_to_known_queues=True).enqueue_job_ids(
                         [str(x) for x in jobs_by_queue[queue]])
 
         print stats

mrq/job.py

@@ -250,7 +250,7 @@ def requeue(self, queue=None, retry_count=0):
             queue = self.data["queue"]
 
         from .queue import Queue
-        queue_obj = Queue(queue)
+        queue_obj = Queue(queue, add_to_known_queues=True)
 
         self._save_status("queued", updates={
             "queue": queue,
@@ -567,7 +567,7 @@ def queue_raw_jobs(queue, params_list, **kwargs):
     """ Queue some jobs on a raw queue """
 
     from .queue import Queue
-    queue_obj = Queue(queue)
+    queue_obj = Queue(queue, add_to_known_queues=True)
     queue_obj.enqueue_raw_jobs(params_list, **kwargs)
 
 
@@ -588,7 +588,7 @@ def queue_jobs(main_task_path, params_list, queue=None, batch_size=1000):
         queue = task_def.get("queue", "default")
 
     from .queue import Queue
-    queue_obj = Queue(queue)
+    queue_obj = Queue(queue, add_to_known_queues=True)
 
     if queue_obj.is_raw:
         raise Exception("Can't queue regular jobs on a raw queue")

mrq/queue.py

@@ -21,7 +21,7 @@ class Queue(object):
     # of Queue in the current process
     known_queues = {}
 
-    def __init__(self, queue_id, add_to_known_queues=True):
+    def __init__(self, queue_id, add_to_known_queues=False):
 
         if isinstance(queue_id, Queue):
             self.id = queue_id.id  # TODO use __new__?
@@ -110,7 +110,7 @@ def redis_known_subqueues(self):
 
         for key in Queue.known_queues:
             if key.startswith(self.id) and not key.endswith(delimiter):
-                queues.append(Queue(key))
+                queues.append(Queue(key, add_to_known_queues=True))
 
         return queues
 
@@ -235,11 +235,35 @@ def all_active(cls):
         return queues
 
     @classmethod
-    def all_known(cls, ):
+    def all_known(cls):
         """ List all previously known queues """
 
-        # raw queues we know exist from the config + known queues in redis
-        return set(context.get_current_config().get("raw_queues", {}).keys() + cls.redis_known_queues().keys())
+        # queues we know exist from the config + known queues in redis
+        return cls.all_known_from_config().union(set(cls.redis_known_queues().keys()))
+
+    @classmethod
+    def all_known_from_config(cls):
+        """ List all known queues from config (raw and regular). Caution: this does not account for the
+            configuration of workers' queues (usually given via command line)
+        """
+
+        cfg = context.get_current_config()
+
+        queues_from_config = [
+          t.get("queue")
+          for t in (cfg.get("tasks") or {}).values()
+          if t.get("queue")
+        ]
+
+        queues_from_config += (cfg.get("raw_queues") or {}).keys()
+
+        queues_from_config += [
+          t.get("retry_queue")
+          for t in (cfg.get("raw_queues") or {}).values()
+          if t.get("retry_queue")
+        ]
+
+        return set(queues_from_config)
 
     @classmethod
     def all(cls):

mrq/worker.py

@@ -73,7 +73,7 @@ def __init__(self):
         self.log_handler = LogHandler(quiet=self.config["quiet"])
         self.log = self.log_handler.get_logger(worker=self.id)
 
-        self.queues = [Queue(x) for x in self.config["queues"] if x]
+        self.queues = [Queue(x, add_to_known_queues=True) for x in self.config["queues"] if x]
 
         self.log.info(
             "Starting Gevent pool with %s worker greenlets (+ report, logs, adminhttp)" %
@@ -212,9 +212,9 @@ def greenlet_subqueues(self):
             try:
                 for queue in self.config["queues"]:
                     if queue.endswith(get_current_config().get("subqueues_delimiter")):
-                        queues += Queue(queue).redis_known_subqueues()
+                        queues += Queue(queue, add_to_known_queues=True).redis_known_subqueues()
                     else:
-                        queues.append(Queue(queue))
+                        queues.append(Queue(queue, add_to_known_queues=True))
 
             except Exception as e:  # pylint: disable=broad-except
                 self.log.error("When refreshing subqueues: %s", e)