refactor queuesize update

Author: maelorn

mrq/basetasks/cleaning.py

@@ -27,6 +27,7 @@ class RequeueRetryJobs(Task):
     max_concurrency = 1
 
     def run(self, params):
+        print("IN")
         return run_task("mrq.basetasks.utils.JobAction", {
             "status": "retry",
             "dateretry": {"$lte": datetime.datetime.utcnow()},

mrq/basetasks/utils.py

@@ -135,22 +135,22 @@ 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:
-                        updates = {
-                            "status": "queued",
-                            "datequeued": datetime.datetime.utcnow(),
-                            "dateupdated": datetime.datetime.utcnow()
-                        }
-
-                        if destination_queue is not None:
-                            updates["queue"] = destination_queue
-
-                        if action == "requeue":
-                            updates["retry_count"] = 0
-
-                        self.collection.update({
-                            "_id": {"$in": jobs_by_queue[queue]}
-                        }, {"$set": updates}, multi=True)
+                for queue in jobs_by_queue:
+                    updates = {
+                        "status": "queued",
+                        "datequeued": datetime.datetime.utcnow(),
+                        "dateupdated": datetime.datetime.utcnow()
+                    }
+
+                    if destination_queue is not None:
+                        updates["queue"] = destination_queue
+
+                    if action == "requeue":
+                        updates["retry_count"] = 0
+
+                    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()})
 

mrq/job.py

@@ -453,14 +453,26 @@ 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" 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)
+            queue = (updates or {}).get("queue") or self.data["queue"]
+            if status != "started":
+                # Queue change
+                if queue != self.data["queue"]:
                     pipe.decr("queuesize:%s" % self.data["queue"])
-                pipe.expire("queuesize:%s" % self.data["queue"], context.get_current_config().get("queue_ttl"))
+                    if status == "queued":
+                        pipe.incr("queuesize:%s" % queue)
+
+                # Regular queues
+                elif status == "queued" and self.data.get("status") != "started":
+                    pipe.incr("queuesize:%s" % queue)
+
+                elif status != "queued" and not self.data.get("raw_queue"):
+                    pipe.decr("queuesize:%s" % queue)
+
+                # Raw queues retries
+                elif (updates or {}).get("retry_count", 0) > self.data.get("retry_count", 0):
+                    pipe.incr("queuesize:%s" % queue)
+
+                pipe.expire("queuesize:%s" % queue, context.get_current_config().get("queue_ttl"))
             pipe.execute()
 
         now = datetime.datetime.utcnow()
@@ -637,7 +649,6 @@ def queue_job(main_task_path, params, **kwargs):
     return queue_jobs(main_task_path, [params], **kwargs)[0]
 
 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:

mrq/queue_raw.py

@@ -227,6 +227,7 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
 
         job_data = [job_factory(p) for p in params]
         for j in job_data:
+            print("J")
             j["status"] = "started"
             j["queue"] = retry_queue
             j["datequeued"] = datetime.now()

tests/test_interrupts.py

@@ -192,7 +192,7 @@ def test_interrupt_worker_sigkill(worker, p_flags):
         "datestarted": datetime.datetime.utcnow() - datetime.timedelta(seconds=300)
     }})
 
-    assert Queue("default").size() == 0
+    assert Queue("default").size() == 1
 
     worker.start(queues="cleaning", deps=False, flush=False,
                  flags=" --config tests/fixtures/config-shorttimeout.py")
@@ -203,7 +203,7 @@ def test_interrupt_worker_sigkill(worker, p_flags):
     assert res["requeued"] == 0
     assert res["started"] == 2  # current job should count too
 
-    assert Queue("default").size() == 0
+    assert Queue("default").size() == 1
 
     job = Job(job_id).fetch().data
     assert job["status"] == "started"