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

* 'master' of github.com:pricingassistant/mrq:
  Added a --max_memory option to restart workers when they go OOM

Author: sylvinus

mrq/config.py

@@ -235,6 +235,14 @@ def add_parser_args(parser, config_type):
             help='Gevent: max number of jobs to do before quitting.' +
                  ' Temp workaround for memory leaks')
 
+        parser.add_argument(
+            '--max_memory',
+            default=0,
+            type=int,
+            action='store',
+            help='Max memory (in Mb) after which the process will be shut down. Use with --processes [1-N]' +
+                 'to have supervisord automatically respawn the worker when this happens')
+
         parser.add_argument(
             '--greenlets',
             '--gevent',  # deprecated

mrq/worker.py

@@ -300,6 +300,10 @@ def report_worker(self, w=0):
 
         report = self.get_worker_report(with_memory=True)
 
+        if self.config["max_memory"] > 0:
+            if report["process"]["mem"]["total"] > (self.config["max_memory"] * 1024 * 1024):
+                self.shutdown_max_memory()
+
         if self.config["report_file"]:
             with open(self.config["report_file"], "wb") as f:
                 f.write(json.dumps(report, ensure_ascii=False))  # pylint: disable=no-member
@@ -398,6 +402,9 @@ def work_loop(self, max_jobs=None):
 
             while True:
 
+                if self.graceful_stop:
+                    break
+
                 while True:
 
                     free_pool_slots = self.gevent_pool.free_count()
@@ -567,6 +574,14 @@ def shutdown_graceful(self):
         self.log.info("Graceful shutdown...")
         raise StopRequested()  # pylint: disable=nonstandard-exception
 
+    def shutdown_max_memory(self):
+
+        # Not in the exitcodes list: we want it to be restarted.
+        self.exitcode = 4
+
+        self.log.info("Max memory reached, shutdown...")
+        self.graceful_stop = True
+
     def shutdown_now(self):
         """ Forced shutdown: interrupts all the jobs. """
 

tests/test_memoryleaks.py

@@ -1,6 +1,27 @@
 import time
 
 
+def test_max_memory_restart(worker):
+
+    N = 20
+
+    worker.start(
+        flags="--processes 1 --greenlets 1 --max_memory 50 --report_interval 1")
+
+    worker.send_tasks(
+        "tests.tasks.general.Leak",
+        [{"size": 1000000, "sleep": 1} for _ in range(N)],
+        queue="default",
+        block=True
+    )
+
+    assert worker.mongodb_jobs.mrq_jobs.find(
+        {"status": "success"}).count() == N
+
+    # We must have been restarted at least once.
+    assert worker.mongodb_jobs.mrq_workers.find().count() > 1
+
+
 def get_diff_after_jobs(worker, n_tasks, leak, sleep=0):
 
     time.sleep(3)
@@ -51,6 +72,8 @@ def test_memoryleaks_noleak(worker):
     assert diff100 < 10000
     assert diff200 < 10000
 
+    assert worker.mongodb_jobs.mrq_workers.find().count() == 1
+
 
 def test_memoryleaks_1mleak(worker):
 
@@ -70,3 +93,5 @@ def test_memoryleaks_1mleak(worker):
 
     assert worker.mongodb_jobs.mrq_jobs.find(
         {"memory_diff": {"$gte": 80000}}).count() == 10
+
+    assert worker.mongodb_jobs.mrq_workers.find().count() == 1