Merge branch 'master' of github.com:pricingassistant/mrq

Author: adipasquale

mrq/scheduler.py

@@ -36,6 +36,7 @@ def sync_tasks(self, tasks):
         # The date part will be discarded in check()
         task["dailytime"] = datetime.datetime.combine(datetime.datetime.utcnow(), task["dailytime"])
         task["interval"] = 3600 * 24
+
       self.collection.insert(task)
       log.debug("Scheduler: added %s" % task["hash"])
 
@@ -55,19 +56,32 @@ def check(self):
 
         dailytime = task.get("dailytime").time()
 
-        if task.get("datelastqueued") and task.get("datelastqueued").time().isoformat()[0:8] != dailytime.isoformat()[0:8]:
-          log.debug("Adjusting the time of scheduled task %s to %s" % (task["_id"], dailytime))
-
-          self.collection.update({"_id": task["_id"]}, {"$set": {
-            "datelastqueued": datetime.datetime.combine(task.get("datelastqueued").date() - datetime.timedelta(days=1), dailytime)
+        time_datelastqueued = task.get("datelastqueued").time().isoformat()[0:8]
+        time_dailytime = dailytime.isoformat()[0:8]
+        if task.get("datelastqueued") and time_datelastqueued != time_dailytime:
+          log.debug("Adjusting the time of scheduled task %s from %s to %s" % (task["_id"], time_datelastqueued, time_dailytime))
+
+          # Make sure we don't queue the task in a loop by adjusting the time
+          if time_datelastqueued < time_dailytime:
+            adjusted_datelastqueued = datetime.datetime.combine(task.get("datelastqueued").date() - datetime.timedelta(days=1), dailytime)
+          else:
+            adjusted_datelastqueued = datetime.datetime.combine(task.get("datelastqueued").date(), dailytime)
+
+          # We do find_and_modify and not update() because several check() may be happening
+          # at the same time.
+          self.collection.find_and_modify({
+            "_id": task["_id"],
+            "datelastqueued": task.get("datelastqueued")
+          }, {"$set": {
+            "datelastqueued": adjusted_datelastqueued
           }})
           self.refresh()
 
       task_data = self.collection.find_and_modify({
         "_id": task["_id"],
         "datelastqueued": {"$lt": last_time}
       }, {"$set": {
-        "datelastqueued": datetime.datetime.utcnow()
+        "datelastqueued": now
       }})
 
       if task_data:

mrq/worker.py

@@ -110,8 +110,8 @@ def connect(self, force=False):
 
   def ensure_indexes(self):
 
-    self.mongodb_logs.mrq_logs.ensure_index([("job", 1)], background=True)
-    self.mongodb_logs.mrq_logs.ensure_index([("worker", 1)], background=True, sparse=True)
+    self.mongodb_logs.mrq_logs.ensure_index([("job", 1)], background=False)
+    self.mongodb_logs.mrq_logs.ensure_index([("worker", 1)], background=False, sparse=True)
 
     if self.config["mongodb_logs_size"] > 0:
 
@@ -120,14 +120,16 @@ def ensure_indexes(self):
       except:
         pass
 
-    self.mongodb_logs.mrq_workers.ensure_index([("status", 1)], background=True)
-    self.mongodb_logs.mrq_workers.ensure_index([("datereported", 1)], background=True, expireAfterSeconds=3600)
+    self.mongodb_logs.mrq_workers.ensure_index([("status", 1)], background=False)
+    self.mongodb_logs.mrq_workers.ensure_index([("datereported", 1)], background=False, expireAfterSeconds=3600)
 
-    self.mongodb_jobs.mrq_jobs.ensure_index([("status", 1)], background=True)
-    self.mongodb_jobs.mrq_jobs.ensure_index([("path", 1), ("status", 1)], background=True)
-    self.mongodb_jobs.mrq_jobs.ensure_index([("worker", 1), ("status", 1)], background=True, sparse=True)
-    self.mongodb_jobs.mrq_jobs.ensure_index([("queue", 1), ("status", 1)], background=True)
-    self.mongodb_jobs.mrq_jobs.ensure_index([("dateexpires", 1)], sparse=True, background=True, expireAfterSeconds=0)
+    self.mongodb_jobs.mrq_jobs.ensure_index([("status", 1)], background=False)
+    self.mongodb_jobs.mrq_jobs.ensure_index([("path", 1), ("status", 1)], background=False)
+    self.mongodb_jobs.mrq_jobs.ensure_index([("worker", 1), ("status", 1)], background=False, sparse=True)
+    self.mongodb_jobs.mrq_jobs.ensure_index([("queue", 1), ("status", 1)], background=False)
+    self.mongodb_jobs.mrq_jobs.ensure_index([("dateexpires", 1)], sparse=True, background=False, expireAfterSeconds=0)
+
+    self.mongodb_jobs.mrq_scheduled_jobs.ensure_index([("hash", 1)], unique=True, background=False, drop_dups=True)
 
     try:
       # This will be default in MongoDB 2.6

tests/fixtures/config-scheduler3.py

@@ -1,12 +1,13 @@
 import datetime
+import os
 
 SCHEDULER_TASKS = [
   {
     "path": "mrq.basetasks.tests.general.MongoInsert",
     "params": {
       "a": 1
     },
-    "dailytime": (datetime.datetime.utcnow() + datetime.timedelta(seconds=3)).time()
+    "dailytime": datetime.datetime.fromtimestamp(float(os.environ.get("MRQ_TEST_SCHEDULER_TIME"))).time()
   },
 ]
 

tests/test_scheduler.py

@@ -1,11 +1,21 @@
 from bson import ObjectId
 import urllib2
 import time
+import pytest
+import datetime
 
 
-def test_scheduler_simple(worker):
+# We want to test that launching the scheduler several times queues tasks only once.
+PROCESS_CONFIGS = [
+  ["--gevent 1"],
+  ["--gevent 1 --processes 5"]
+]
 
-  worker.start(flags="--scheduler --config tests/fixtures/config-scheduler1.py")
+
+@pytest.mark.parametrize(["p_flags"], PROCESS_CONFIGS)
+def test_scheduler_simple(worker, p_flags):
+
+  worker.start(flags="--scheduler --config tests/fixtures/config-scheduler1.py %s" % p_flags)
 
   collection = worker.mongodb_logs.tests_inserts
   scheduled_jobs = worker.mongodb_jobs.mrq_scheduled_jobs
@@ -30,7 +40,7 @@ def test_scheduler_simple(worker):
   collection.remove({})
 
   # Start with new config
-  worker.start(deps=False, flags="--scheduler --config tests/fixtures/config-scheduler2.py")
+  worker.start(deps=False, flags="--scheduler --config tests/fixtures/config-scheduler2.py %s" % p_flags)
 
   time.sleep(2)
 
@@ -43,14 +53,21 @@ def test_scheduler_simple(worker):
   assert len(inserts) == 3
 
 
-def test_scheduler_dailytime(worker):
+@pytest.mark.parametrize(["p_flags"], PROCESS_CONFIGS)
+def test_scheduler_dailytime(worker, p_flags):
 
   # Task is scheduled in 3 seconds
-  worker.start(flags="--scheduler --config tests/fixtures/config-scheduler3.py")
+  worker.start(
+    flags="--scheduler --config tests/fixtures/config-scheduler3.py %s" % p_flags,
+    env={
+
+      # We need to pass this in the environment so that each worker has the exact same hash
+      "MRQ_TEST_SCHEDULER_TIME": str(time.time() + 4)
+    })
 
   # It will be done a first time immediately
 
-  time.sleep(1)
+  time.sleep(3)
 
   collection = worker.mongodb_logs.tests_inserts
 
@@ -60,3 +77,9 @@ def test_scheduler_dailytime(worker):
   time.sleep(4)
 
   assert collection.find().count() == 2
+
+  # Nothing more should happen today
+  time.sleep(4)
+
+  assert collection.find().count() == 2
+