Fix tests & restore queued order

Author: sylvinus

mrq/basetasks/utils.py

@@ -120,6 +120,7 @@ def perform_action(self, action, query, destination_queue):
 
                     updates = {
                         "status": "queued",
+                        "datequeued": datetime.datetime.utcnow(),
                         "dateupdated": datetime.datetime.utcnow()
                     }
 

mrq/job.py

@@ -266,6 +266,7 @@ def requeue(self, queue=None, retry_count=0):
 
         self._save_status("queued", updates={
             "queue": queue,
+            "datequeued": datetime.datetime.utcnow(),
             "retry_count": retry_count
         })
 
@@ -632,6 +633,7 @@ def queue_jobs(main_task_path, params_list, queue=None, batch_size=1000):
             "path": main_task_path,
             "params": params,
             "queue": queue,
+            "datequeued": datetime.datetime.utcnow(),
             "status": "queued"
         } for params in params_group], w=1, return_jobs=False)
 

mrq/monkey.py

@@ -389,10 +389,6 @@ class mrq_patched_pymongo_cursor(Cursor):
         # Some dark magic is needed here to cope with python's name mangling for private variables.
         def _Cursor__send_message(self, *args, **kwargs):
 
-            if len(args) and args[0].name == "find":
-                return Cursor._Cursor__send_message(self, *args, **kwargs)
-
-            # print self.__dict__
             job = get_current_job()
 
             if job:

mrq/queue_raw.py

@@ -142,7 +142,6 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
         retry_queue = self.get_retry_queue()
 
         params = []
-        jobs = []
 
         # ZSET with times
         if self.is_timed:
@@ -187,10 +186,11 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
             params = redis_group_command("lpop", max_jobs, self.redis_key)
 
         if len(params) == 0:
-            return []
+            return
 
         if worker:
             worker.status = "spawn"
+            worker.idle_event.clear()
 
         job_data = [job_factory(p) for p in params]
         for j in job_data:
@@ -200,9 +200,8 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
             if worker:
                 j["worker"] = worker.id
 
-        jobs += job_class.insert(job_data, statuses_no_storage=statuses_no_storage)
-
-        return jobs
+        for job in job_class.insert(job_data, statuses_no_storage=statuses_no_storage):
+            yield job
 
     def get_sorted_graph(
             self,

mrq/queue_regular.py

@@ -18,7 +18,7 @@ def size(self):
     def list_job_ids(self, skip=0, limit=20):
         """ Returns a list of job ids on a queue """
 
-        return [str(x) for x in self.collection.find(
+        return [str(x["_id"]) for x in self.collection.find(
             {"status": "queued"},
             sort=[("_id", -1 if self.is_reverse else 1)],
             projection={"_id": 1})
@@ -31,25 +31,49 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
             from .job import Job
             job_class = Job
 
-        if worker:
-            worker.status = "spawn"
-            worker.idle_event.clear()
-
         count = 0
 
-        for _ in range(max_jobs):
+        job_ids = None
 
-            job_data = self.collection.find_one_and_update(
+        # TODO: remove _id sort after full migration to datequeued
+        sort_order = [("datequeued", -1 if self.is_reverse else 1), ("_id", -1 if self.is_reverse else 1)]
+
+        # MongoDB optimization: with many jobs it's faster to fetch the IDs first and do the atomic update second
+        # Some jobs may have been stolen by another worker in the meantime but it's a balance (should we over-fetch?)
+        if max_jobs > 5:
+            job_ids = [x["_id"] for x in self.collection.find(
                 {
                     "status": "queued",
                     "queue": self.id
                 },
+                limit=max_jobs,
+                sort=sort_order,
+                projection={"_id": 1}
+            )]
+
+            if len(job_ids) == 0:
+                return
+
+        for i in range(max_jobs if job_ids is None else len(job_ids)):
+
+            query = {
+                "status": "queued",
+                "queue": self.id
+            }
+            if job_ids is not None:
+                query = {
+                    "status": "queued",
+                    "_id": job_ids[i]
+                }
+
+            job_data = self.collection.find_one_and_update(
+                query,
                 {"$set": {
                     "status": "started",
                     "datestarted": datetime.datetime.utcnow(),
                     "worker": worker.id if worker else None
                 }},
-                sort=[("_id", -1 if self.is_reverse else 1)],
+                sort=sort_order,
                 return_document=ReturnDocument.AFTER,
                 projection={
                     "_id": 1,
@@ -64,6 +88,10 @@ def dequeue_jobs(self, max_jobs=1, job_class=None, worker=None):
             if not job_data:
                 break
 
+            if worker:
+                worker.status = "spawn"
+                worker.idle_event.clear()
+
             count += 1
             context.metric("queues.%s.dequeued" % job_data["queue"], 1)
 

mrq/worker.py

@@ -144,6 +144,8 @@ def ensure_indexes(self):
             [("dateexpires", 1)], sparse=True, background=False, expireAfterSeconds=0)
         self.mongodb_jobs.mrq_jobs.ensure_index(
             [("dateretry", 1)], sparse=True, background=False)
+        self.mongodb_jobs.mrq_jobs.ensure_index(
+            [("datequeued", 1)], sparse=True, background=True)
 
         self.mongodb_jobs.mrq_scheduled_jobs.ensure_index(
             [("hash", 1)], unique=True, background=False, drop_dups=True)

tests/test_general.py

@@ -112,6 +112,26 @@ def test_general_simple_task_multiple(worker):
         [["dateupdated", 1]])] == [42, 42, 41]
 
 
+def test_general_requeue_order(worker):
+    from mrq.job import Job
+
+    jobids = worker.send_tasks("tests.tasks.general.Add", [
+        {"a": 41, "b": 1, "sleep": 4},
+        {"a": 42, "b": 1, "sleep": 1},
+        {"a": 43, "b": 1, "sleep": 1}
+    ], block=False)
+
+    time.sleep(2)
+
+    # We should be executing job1 now. Let's requeue job2, making it go to the end of the queue.
+    Job(jobids[1]).requeue()
+
+    worker.get_wait_for_idle()
+
+    assert [x["result"] for x in worker.mongodb_jobs.mrq_jobs.find().sort(
+        [["dateupdated", 1]])] == [42, 44, 43]
+
+
 def test_general_simple_task_reverse(worker):
 
     worker.start(queues="default_reverse xtest test_timed_set", flags="--config tests/fixtures/config-raw1.py")

tests/test_io_hooks.py

@@ -119,7 +119,7 @@ def test_io_hooks_mongodb(worker):
     assert job_events[4]["method"] == "cursor"
     assert job_events[5]["method"] == "cursor"
 
-    # Then getmore query
+    # Then getmore query (can't understand why there are 2 more of those)
     assert job_events[6]["hook"] == "mongodb_pre"
     assert job_events[7]["hook"] == "mongodb_post"
 
@@ -129,15 +129,25 @@ def test_io_hooks_mongodb(worker):
     assert job_events[6]["method"] == "cursor"
     assert job_events[7]["method"] == "cursor"
 
-    # Result MongoDB update
-
+    # Then getmore query
     assert job_events[8]["hook"] == "mongodb_pre"
     assert job_events[9]["hook"] == "mongodb_post"
 
-    assert job_events[8]["method"] == "update"
-    assert job_events[9]["method"] == "update"
+    assert job_events[8]["collection"] == "mrq.tests_inserts"
+    assert job_events[9]["collection"] == "mrq.tests_inserts"
+
+    assert job_events[8]["method"] == "cursor"
+    assert job_events[9]["method"] == "cursor"
+
+    # Result MongoDB update
+
+    assert job_events[10]["hook"] == "mongodb_pre"
+    assert job_events[11]["hook"] == "mongodb_post"
+
+    assert job_events[10]["method"] == "update"
+    assert job_events[11]["method"] == "update"
 
-    assert job_events[8]["collection"] == "mrq.mrq_jobs"
-    assert job_events[9]["collection"] == "mrq.mrq_jobs"
+    assert job_events[10]["collection"] == "mrq.mrq_jobs"
+    assert job_events[11]["collection"] == "mrq.mrq_jobs"
 
-    assert len(job_events) == 5 * 2
+    assert len(job_events) == 6 * 2

tests/test_jobinspect.py

@@ -102,9 +102,10 @@ def test_current_job_trace_io(worker, p_testtype, p_testparams, p_type, p_data,
                     admin_worker = {}
                 if len(admin_worker.get("jobs", [])) > 0:
                     io = admin_worker["jobs"][0].get("io")
+
                     # Don't take MRQ's IOs as regular IO
                     if io:
-                        if io["type"] == "mongodb" and io["data"]["collection"] in ["mrq.mrq_jobs", "mrq.mrq_logs"]:
+                        if io["type"].startswith("mongodb") and io["data"]["collection"] in ["mrq.mrq_jobs", "mrq.mrq_logs"]:
                             io = False
                         else:
                             break