fix tests

Author: maelorn

mrq/basetasks/cleaning.py

@@ -57,7 +57,7 @@ def run(self, params):
         # iterate over them.
 
         fields = {
-            "_id": 1, "datestarted": 1, "queue": 1, "path": 1, "retry_count": 1, "worker": 1
+            "_id": 1, "datestarted": 1, "queue": 1, "path": 1, "retry_count": 1, "worker": 1, "status": 1
         }
         for job_data in connections.mongodb_jobs.mrq_jobs.find(
                 {"status": "started"}, projection=fields):

mrq/basetasks/utils.py

@@ -5,6 +5,7 @@
 from mrq.queue import Queue
 from bson import ObjectId
 from mrq.context import connections, get_current_config, get_current_job
+from mrq.job import set_queues_size
 from collections import defaultdict
 from mrq.utils import group_iter
 import datetime
@@ -59,7 +60,7 @@ def build_query(self):
 
             for key in params_dict:
                 query["params.%s" % key] = params_dict[key]
-                
+
         if current_job and "_id" not in query:
             query["_id"] = {"$lte": current_job.id}
 
@@ -91,13 +92,24 @@ def perform_action(self, action, query, destination_queue):
                 result_ttl = max([default_job_timeout] + tasks_ttls)
 
             now = datetime.datetime.utcnow()
+
+            size_by_queues = defaultdict(int)
+            if "queue" not in query:
+                for job in self.collection.find(query, projection={"queue": 1}):
+                    size_by_queues[job["queue"]] += 1
+
             ret = self.collection.update(query, {"$set": {
                 "status": "cancel",
                 "dateexpires": now + datetime.timedelta(seconds=result_ttl),
                 "dateupdated": now
             }}, multi=True)
             stats["cancelled"] = ret["n"]
 
+            if "queue" in query:
+                if isinstance(query["queue"], str):
+                    size_by_queues[query["queue"]] = ret["n"]
+            set_queues_size(size_by_queues, action="decr")
+
             # Special case when emptying just by queue name: empty it directly!
             # In this case we could also loose some jobs that were queued after
             # the MongoDB update. They will be "lost" and requeued later like the other case
@@ -123,22 +135,23 @@ def perform_action(self, action, query, destination_queue):
                     jobs_by_queue[job["queue"]].append(job["_id"])
                     stats["requeued"] += 1
 
-                for queue in jobs_by_queue:
+                    for queue in jobs_by_queue:
+                        updates = {
+                            "status": "queued",
+                            "datequeued": datetime.datetime.utcnow(),
+                            "dateupdated": datetime.datetime.utcnow()
+                        }
 
-                    updates = {
-                        "status": "queued",
-                        "datequeued": datetime.datetime.utcnow(),
-                        "dateupdated": datetime.datetime.utcnow()
-                    }
+                        if destination_queue is not None:
+                            updates["queue"] = destination_queue
 
-                    if destination_queue is not None:
-                        updates["queue"] = destination_queue
+                        if action == "requeue":
+                            updates["retry_count"] = 0
 
-                    if action == "requeue":
-                        updates["retry_count"] = 0
+                        self.collection.update({
+                            "_id": {"$in": jobs_by_queue[queue]}
+                        }, {"$set": updates}, multi=True)
 
-                    self.collection.update({
-                        "_id": {"$in": jobs_by_queue[queue]}
-                    }, {"$set": updates}, multi=True)
+                set_queues_size({queue: len(jobs) for queue, jobs in jobs_by_queue.iteritems()})
 
         return stats

mrq/job.py

@@ -453,10 +453,12 @@ def _save_status(self, status, updates=None, exception=False, w=None, j=None):
             return
 
         with context.connections.redis.pipeline(transaction=False) as pipe:
-            if status != "started":
-                if status == "queued":
+            if status != "started" and "raw_queue" not in self.data:
+                if status == "queued" and self.data["status"] != "started":
+                    print("INCR", self.data["queue"])
                     pipe.incr("queuesize:%s" % self.data["queue"])
                 else:
+                    print("DECR", self.data["queue"], self.data)
                     pipe.decr("queuesize:%s" % self.data["queue"])
                 pipe.expire("queuesize:%s" % self.data["queue"], context.get_current_config().get("queue_ttl"))
             pipe.execute()
@@ -634,8 +636,14 @@ def queue_job(main_task_path, params, **kwargs):
 
     return queue_jobs(main_task_path, [params], **kwargs)[0]
 
-def set_queue_size(queue, size):
-    context.connections.redis.setex("queuesize:%s" % queue, context.get_current_config().get("queue_ttl"), size)
+def set_queues_size(size_by_queues, action="incr"):
+    print("BLA", size_by_queues, action)
+    if len(size_by_queues) > 0:
+        with context.connections.redis.pipeline(transaction=False) as pipe:
+            for queue in size_by_queues:
+                getattr(pipe, action)("queuesize:%s" % queue, amount=size_by_queues[queue])
+                pipe.expire("queuesize:%s" % queue, context.get_current_config().get("queue_ttl"))
+            pipe.execute()
 
 def queue_jobs(main_task_path, params_list, queue=None, batch_size=1000):
     """ Queue multiple jobs on a regular queue """
@@ -669,6 +677,6 @@ def queue_jobs(main_task_path, params_list, queue=None, batch_size=1000):
         all_ids += job_ids
 
     queue_obj.notify(len(all_ids))
-    set_queue_size(queue, len(all_ids))
+    set_queues_size({queue: len(all_ids)})
 
     return all_ids