Merge branch '0.9.x' into task_expiry

* 0.9.x:
  more readable condition
  fix cleaning logic
  make sure there's no job before cleaning a known queue
  Store subpool tracebacks + test fixes
  Fix Python 3 incompatibility (#188)

Author: ismael

mrq/basetasks/cleaning.py

@@ -129,16 +129,19 @@ def run(self, params):
 
         # Only clean queues older than N days
         time_threshold = time.time() - max_age
-        for queue, time_last_used in known_queues.iteritems():
+        for queue, time_last_used in known_queues.items():
             if queue in queues_from_config:
                 continue
             if time_last_used < time_threshold:
                 q = Queue(queue, add_to_known_queues=False)
+                # size() returns the number of queued jobs
                 size = q.size()
+
+                has_job = None
                 if check_mongo:
                     has_job = connections.mongodb_jobs.mrq_jobs.find_one({"queue": queue})
 
-                if size == 0 or has_job is None:
+                if size == 0 and has_job is None:
                     removed_queues.append(queue)
                     print("Removing empty queue '%s' from known queues ..." % queue)
                     if not pretend:

mrq/context.py

@@ -10,8 +10,9 @@
 import pymongo
 import traceback
 from .utils import LazyObject, load_class_by_path
-from itertools import count as itertools_count
 from .config import get_config
+from .subpool import subpool_map, subpool_imap
+
 
 # This should be MRQ's only Python object shared by all the jobs in the same process
 _GLOBAL_CONTEXT = {
@@ -202,121 +203,6 @@ def trace(*args):
     greenlet.settrace(trace)  # pylint: disable=no-member
 
 
-def subpool_map(pool_size, func, iterable):
-    """ Starts a Gevent pool and run a map. Takes care of setting current_job and cleaning up. """
-
-    if not pool_size:
-        return [func(*args) for args in iterable]
-
-    counter = itertools_count()
-
-    current_job = get_current_job()
-
-    def inner_func(*args):
-        """ As each call to 'func' will be done in a random greenlet of the subpool, we need to
-            register their IDs with set_current_job() to make get_current_job() calls work properly
-            inside 'func'.
-        """
-        next(counter)
-        if current_job:
-            set_current_job(current_job)
-
-        try:
-          ret = func(*args)
-        except Exception as exc:
-          trace = traceback.format_exc()
-          log.error("Error in subpool: %s \n%s" % (exc, trace))
-          raise
-
-        if current_job:
-            set_current_job(None)
-        return ret
-
-    def inner_iterable():
-        """ This will be called inside the pool's main greenlet, which ID also needs to be registered """
-        if current_job:
-            set_current_job(current_job)
-
-        for x in iterable:
-            yield x
-
-        if current_job:
-            set_current_job(None)
-
-    start_time = time.time()
-    pool = gevent.pool.Pool(size=pool_size)
-    ret = pool.map(inner_func, inner_iterable())
-    pool.join(raise_error=True)
-    total_time = time.time() - start_time
-
-    log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))
-
-    return ret
-
-
-def subpool_imap(pool_size, func, iterable, flatten=False, unordered=False, buffer_size=None):
-  """ Generator version of subpool_map. Should be used with unordered=True for optimal performance """
-
-  if not pool_size:
-    for args in iterable:
-      yield func(*args)
-
-  counter = itertools_count()
-
-  current_job = get_current_job()
-
-  def inner_func(*args):
-    """ As each call to 'func' will be done in a random greenlet of the subpool, we need to
-        register their IDs with set_current_job() to make get_current_job() calls work properly
-        inside 'func'.
-    """
-    next(counter)
-    if current_job:
-      set_current_job(current_job)
-
-    try:
-      ret = func(*args)
-    except Exception as exc:
-      trace = traceback.format_exc()
-      log.error("Error in subpool: %s \n%s" % (exc, trace))
-      raise
-
-    if current_job:
-      set_current_job(None)
-    return ret
-
-  def inner_iterable():
-    """ This will be called inside the pool's main greenlet, which ID also needs to be registered """
-    if current_job:
-      set_current_job(current_job)
-
-    for x in iterable:
-      yield x
-
-    if current_job:
-      set_current_job(None)
-
-  start_time = time.time()
-  pool = gevent.pool.Pool(size=pool_size)
-
-  if unordered:
-    iterator = pool.imap_unordered(inner_func, inner_iterable(), maxsize=buffer_size or pool_size)
-  else:
-    iterator = pool.imap(inner_func, inner_iterable())
-
-  for x in iterator:
-    if flatten:
-      for y in x:
-        yield y
-    else:
-      yield x
-
-  pool.join(raise_error=True)
-  total_time = time.time() - start_time
-
-  log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))
-
-
 def run_task(path, params):
     """ Runs a task code synchronously """
     task_class = load_class_by_path(path)

mrq/job.py

@@ -473,8 +473,10 @@ def _save_status(self, status, updates=None, exception=False, w=None, j=None):
         if exception:
             trace = traceback.format_exc()
             context.log.error(trace)
+            exc, value = sys.exc_info()[0:2]
+            if hasattr(value, "subpool_traceback"):
+                trace = "Exception first caught in a subpool. Traceback:\n%s\n%s" % (value.subpool_traceback, trace)
             db_updates["traceback"] = trace
-            exc = sys.exc_info()[0]
             db_updates["exceptiontype"] = exc.__name__
 
             self._save_traceback_history(status, trace, exc)

mrq/subpool.py

@@ -0,0 +1,123 @@
+from itertools import count as itertools_count
+import traceback
+import time
+import gevent
+
+
+def subpool_map(pool_size, func, iterable):
+    """ Starts a Gevent pool and run a map. Takes care of setting current_job and cleaning up. """
+
+    from .context import get_current_job, set_current_job, log
+
+    if not pool_size:
+        return [func(*args) for args in iterable]
+
+    counter = itertools_count()
+
+    current_job = get_current_job()
+
+    def inner_func(*args):
+        """ As each call to 'func' will be done in a random greenlet of the subpool, we need to
+            register their IDs with set_current_job() to make get_current_job() calls work properly
+            inside 'func'.
+        """
+        next(counter)
+        if current_job:
+            set_current_job(current_job)
+
+        try:
+          ret = func(*args)
+        except Exception as exc:
+          trace = traceback.format_exc()
+          exc.subpool_traceback = trace
+          raise
+
+        if current_job:
+            set_current_job(None)
+        return ret
+
+    def inner_iterable():
+        """ This will be called inside the pool's main greenlet, which ID also needs to be registered """
+        if current_job:
+            set_current_job(current_job)
+
+        for x in iterable:
+            yield x
+
+        if current_job:
+            set_current_job(None)
+
+    start_time = time.time()
+    pool = gevent.pool.Pool(size=pool_size)
+    ret = pool.map(inner_func, inner_iterable())
+    pool.join(raise_error=True)
+    total_time = time.time() - start_time
+
+    log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))
+
+    return ret
+
+
+def subpool_imap(pool_size, func, iterable, flatten=False, unordered=False, buffer_size=None):
+  """ Generator version of subpool_map. Should be used with unordered=True for optimal performance """
+
+  from .context import get_current_job, set_current_job, log
+
+  if not pool_size:
+    for args in iterable:
+      yield func(*args)
+
+  counter = itertools_count()
+
+  current_job = get_current_job()
+
+  def inner_func(*args):
+    """ As each call to 'func' will be done in a random greenlet of the subpool, we need to
+        register their IDs with set_current_job() to make get_current_job() calls work properly
+        inside 'func'.
+    """
+    next(counter)
+    if current_job:
+      set_current_job(current_job)
+
+    try:
+      ret = func(*args)
+    except Exception as exc:
+      trace = traceback.format_exc()
+      exc.subpool_traceback = trace
+      raise
+
+    if current_job:
+      set_current_job(None)
+    return ret
+
+  def inner_iterable():
+    """ This will be called inside the pool's main greenlet, which ID also needs to be registered """
+    if current_job:
+      set_current_job(current_job)
+
+    for x in iterable:
+      yield x
+
+    if current_job:
+      set_current_job(None)
+
+  start_time = time.time()
+  pool = gevent.pool.Pool(size=pool_size)
+
+  if unordered:
+    iterator = pool.imap_unordered(inner_func, inner_iterable(), maxsize=buffer_size or pool_size)
+  else:
+    iterator = pool.imap(inner_func, inner_iterable())
+
+  for x in iterator:
+    if flatten:
+      for y in x:
+        yield y
+    else:
+      yield x
+
+  pool.join(raise_error=True)
+  total_time = time.time() - start_time
+
+  log.debug("SubPool ran %s greenlets in %0.6fs" % (counter, total_time))

tests/conftest.py

@@ -261,7 +261,7 @@ def get_report(self, with_memory=False):
         f.close()
         return data
 
-    def wait_for_idle(self):
+    def wait_for_idle(self, timeout=None):
 
         if "--processes" in self.cmdline:
             print("Warning: wait_for_idle() doesn't support multiprocess workers yet")
@@ -275,7 +275,7 @@ def wait_for_idle(self):
             return True
 
         try:
-            wait_for_net_service("127.0.0.1", self.admin_port, poll_interval=0.01)
+            wait_for_net_service("127.0.0.1", self.admin_port, poll_interval=0.01, timeout=timeout)
             f = urllib.request.urlopen("http://127.0.0.1:%s/wait_for_idle" % self.admin_port)
             data = f.read().decode('utf-8')
             assert data == "idle"

tests/tasks/general.py

@@ -3,7 +3,7 @@
 from builtins import range
 from mrq.task import Task
 from mrq.context import (log, retry_current_job, connections, get_current_config, get_current_job,
-                         subpool_map, abort_current_job, set_current_job_progress)
+                         subpool_map, subpool_imap, abort_current_job, set_current_job_progress)
 from mrq.job import queue_job
 import urllib.request, urllib.error, urllib.parse
 import json
@@ -244,7 +244,7 @@ def inner(self, x):
             return True
 
         if x == "exception":
-            raise Exception(x)
+            raise Exception(x)  # __INNER_EXCEPTION_LINE__
 
         time.sleep(x)
 
@@ -253,7 +253,10 @@ def inner(self, x):
     def run(self, params):
         self.job = get_current_job()
 
-        return subpool_map(params["pool_size"], self.inner, params["inner_params"])
+        if params.get("imap"):
+            return subpool_map(params["pool_size"], self.inner, params["inner_params"])
+        else:
+            return list(subpool_imap(params["pool_size"], self.inner, params["inner_params"]))
 
 
 class GetMetrics(Task):

tests/test_agent.py

@@ -11,8 +11,9 @@ def scenario(profiles, agents):
 
     for agent in agents:
         agent["worker_group"] = "xx"
+        agent["status"] = "started"
 
-    connections.mongodb_jobs.mrq_agents.insert_many(agents + [{"worker_group": "zz"}])
+    connections.mongodb_jobs.mrq_agents.insert_many(agents + [{"worker_group": "yy", "status": "started"}, {"worker_group": "zz"}])
     connections.mongodb_jobs.mrq_workergroups.insert_one({"_id": "xx", "profiles": profiles})
 
     agent = Agent(worker_group="xx")

tests/test_general.py

@@ -152,6 +152,8 @@ def test_general_simple_task_reverse(worker):
 def test_known_queues_lifecycle(worker):
 
     worker.start(queues="default_reverse xtest test_timed_set", flags="--config tests/fixtures/config-raw1.py")
+    time.sleep(1)
+    worker.wait_for_idle()
 
     # Test known queues
     from mrq.queue import Queue, send_task

tests/test_parallel.py

@@ -60,7 +60,7 @@ def test_dequeue_strategy(worker, p_greenlets, p_strategy):
     print("Worker has flags %s" % flags)
     worker.start(flags=flags, queues="q1 q2", deps=False, block=False)
 
-    gotit = worker.wait_for_idle()
+    gotit = worker.wait_for_idle(timeout=10)
 
     if p_strategy == "burst":
         assert not gotit  # because worker should be stopped already

tests/test_subpool.py

@@ -49,13 +49,22 @@ def run(params):
         assert total_time < 2
 
 
-def test_subpool_exception(worker):
+@pytest.mark.parametrize(["p_imap"], [
+    [True],
+    [False]
+])
+def test_subpool_exception(worker, p_imap):
 
-    # Exception
+    # An exception in the subpool is raised outside the pool
     worker.send_task("tests.tasks.general.SubPool", {
-        "pool_size": 20, "inner_params": ["exception"]
+        "pool_size": 20, "inner_params": ["exception"], "imap": p_imap
     }, accept_statuses=["failed"])
 
+    job = worker.mongodb_jobs.mrq_jobs.find_one()
+    assert job
+    assert job["status"] == "failed"
+    assert "__INNER_EXCEPTION_LINE__" in job["traceback"]
+
 
 @pytest.mark.parametrize(["p_size"], [
     [0],