Add flags to check memory usage and track leaks. Move mrq_workers to mongodb_jobs

Author: sylvinus

.gitignore

@@ -38,4 +38,6 @@ venv
 pypy
 .DS_Store
 
-mrq-config.py
+mrq-config.py
+dump.rdb
+supervisor.pid

Dockerfile

@@ -6,7 +6,7 @@ RUN echo 'deb http://downloads-distro.mongodb.org/repo/ubuntu-upstart dist 10gen
 RUN apt-get update && echo "Updated on 2014-01-15"
 RUN apt-get upgrade -y
 
-RUN apt-get install -y gcc make g++ build-essential libc6-dev tcl curl adduser mongodb-10gen python python-pip python-dev strace git software-properties-common libev-dev nginx
+RUN apt-get install -y gcc make g++ build-essential libc6-dev tcl curl adduser mongodb-10gen python python-pip python-dev strace git software-properties-common libev-dev nginx graphviz
 
 # Then add PPAs (after software-properties-common is installed)
 # RUN add-apt-repository -y ppa:pypy/ppa

mrq/basetasks/tests/general.py

@@ -4,6 +4,7 @@
 import urllib2
 
 
+
 class Add(Task):
 
   def run(self, params):
@@ -36,6 +37,21 @@ def run(self, params):
     return len(t)
 
 
+LEAKS = []
+
+
+class Leak(Task):
+  def run(self, params):
+
+    if params.get("size", 0) > 0:
+      LEAKS.append("x" * params.get("size", 0))
+
+    if params.get("sleep", 0) > 0:
+      sleep(params.get("sleep", 0))
+
+    return params.get("return")
+
+
 class Retry(Task):
 
   def run(self, params):

mrq/config.py

@@ -29,17 +29,26 @@ def get_config(sources=("file", "env", "args"), env_prefix="MRQ_", defaults=None
   parser.add_argument('--trace_greenlets', action='store_true', default=False,
                       help='Collect stats about each greenlet execution time and switches.')
 
+  parser.add_argument('--trace_memory', action='store_true', default=False,
+                      help='Collect stats about memory for each task. Incompatible with gevent > 1')
+
+  parser.add_argument('--trace_memory_type', action='store', default="",
+                      help='Create a .png object graph in trace_memory_output_dir with a random object of this type.')
+
+  parser.add_argument('--trace_memory_output_dir', action='store', default="memory_traces",
+                      help='Directory where to output .pngs with object graphs')
+
   parser.add_argument('--profile', action='store_true', default=False,
                       help='Run profiling on the whole worker')
 
   parser.add_argument('--mongodb_jobs', action='store', default="mongodb://127.0.0.1:27017/mrq",
-                      help='MongoDB URI for the jobs database')
+                      help='MongoDB URI for the jobs, scheduled_jobs & workers database')
 
   parser.add_argument('--mongodb_logs', action='store', default="mongodb://127.0.0.1:27017/mrq",
-                      help='MongoDB URI for the logs database')
+                      help='MongoDB URI for the logs database. If set to "0", will disable remote logs.')
 
   parser.add_argument('--mongodb_logs_size', action='store', default=16 * 1024 * 1024, type=int,
-                      help='If provided, sets the log collection to capped to that amount of bytes')
+                      help='If provided, sets the log collection to capped to that amount of bytes.')
 
   parser.add_argument('--redis', action='store', default="redis://127.0.0.1:6379",
                       help='Redis URI')

mrq/dashboard/app.py

@@ -91,7 +91,7 @@ def api_datatables(unit):
   if unit == "workers":
     fields = None
     query = {"status": {"$nin": ["stop"]}}
-    collection = connections.mongodb_logs.mrq_workers
+    collection = connections.mongodb_jobs.mrq_workers
     sort = [("datestarted", -1)]
 
     if request.args.get("showstopped"):

mrq/job.py

@@ -7,6 +7,7 @@
 from .queue import Queue
 from .context import get_current_worker, log, connections, get_current_config
 import gevent
+import gc
 
 
 class Job(object):
@@ -232,3 +233,46 @@ def wait(self, poll_interval=1, timeout=None, full_data=False):
       time.sleep(poll_interval)
 
     raise Exception("Waited for job result for %ss seconds, timeout." % timeout)
+
+  def trace_memory_start(self):
+    """ Starts measuring memory consumption """
+    import objgraph
+    objgraph.show_growth(limit=10)
+    #gc.collect()  # - is done in show_growth
+    self._memory_start = self.worker.get_memory()
+
+  def trace_memory_stop(self):
+    """ Stops measuring memory consumption """
+
+    #gc.collect() - is done in show_growth
+    import objgraph, random, gevent
+
+    gevent.sleep(0)
+
+    objgraph.show_growth(limit=10)
+
+    trace_type = get_current_config()["trace_memory_type"]
+    if trace_type:
+
+      objgraph.show_chain(
+        objgraph.find_backref_chain(
+          random.choice(objgraph.by_type(trace_type)),
+          objgraph.is_proper_module
+        ),
+        filename='%s/%s-%s.png' % (get_current_config()["trace_memory_output_dir"], trace_type, self.id)
+      )
+
+    self._memory_stop = self.worker.get_memory()
+
+    diff = self._memory_stop - self._memory_start
+
+    log.debug("Memory diff for job %s : %s" % (self.id, diff))
+
+    # We need to update it later than the results, we need them off memory already.
+    self.collection.update({
+      "_id": self.id
+    }, {"$set": {
+      "memory_diff": diff
+    }}, w=1)
+
+

mrq/logger.py

@@ -42,6 +42,9 @@ def log(self, level, *args, **kwargs):
     if not self.quiet:
       print formatted
 
+    if self.collection is False:
+      return
+
     if worker is not None:
       self.buffer["workers"][worker].append(formatted)
     elif job is not None:
@@ -55,7 +58,7 @@ def log(self, level, *args, **kwargs):
   def flush(self, w=0):
 
     # We may log some stuff before we are even connected to Mongo!
-    if self.collection is None:
+    if not self.collection:
       return
 
     inserts = [{

mrq/worker.py

@@ -44,6 +44,10 @@ class Worker(object):
   """
   status = "init"
 
+  mongodb_jobs = None
+  mongodb_logs = None
+  redis = None
+
   def __init__(self, config):
 
     self.config = config
@@ -99,9 +103,12 @@ def connect(self, force=False):
     # Accessing connections attributes will automatically connect
     self.redis = connections.redis
     self.mongodb_jobs = connections.mongodb_jobs
-    self.mongodb_logs = connections.mongodb_logs
 
-    self.log_handler.set_collection(self.mongodb_logs.mrq_logs)
+    if self.config["mongodb_logs"] == "0":
+      self.log_handler.set_collection(False)  # Disable
+    else:
+      self.mongodb_logs = connections.mongodb_logs
+      self.log_handler.set_collection(self.mongodb_logs.mrq_logs)
 
     self.connected = True
 
@@ -110,18 +117,20 @@ def connect(self, force=False):
 
   def ensure_indexes(self):
 
-    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.mongodb_logs:
 
-    if self.config["mongodb_logs_size"] > 0:
+      self.mongodb_logs.mrq_logs.ensure_index([("job", 1)], background=False)
+      self.mongodb_logs.mrq_logs.ensure_index([("worker", 1)], background=False, sparse=True)
 
-      try:
-        self.mongodb_logs.command("convertToCapped", "mrq_logs", size=self.config["mongodb_logs_size"])
-      except:
-        pass
+      if self.config["mongodb_logs_size"] > 0:
+
+        try:
+          self.mongodb_logs.command("convertToCapped", "mrq_logs", size=self.config["mongodb_logs_size"])
+        except:
+          pass
 
-    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_workers.ensure_index([("status", 1)], background=False)
+    self.mongodb_jobs.mrq_workers.ensure_index([("datereported", 1)], background=False, expireAfterSeconds=3600)
 
     self.mongodb_jobs.mrq_jobs.ensure_index([("status", 1)], background=False)
     self.mongodb_jobs.mrq_jobs.ensure_index([("path", 1), ("status", 1)], background=False)
@@ -166,6 +175,9 @@ def greenlet_monitoring(self):
       self.flush_logs(w=0)
       time.sleep(int(self.config["report_interval"]))
 
+  def get_memory(self):
+    return self.process.get_memory_info().rss
+
   def get_worker_report(self):
 
     greenlets = []
@@ -217,7 +229,7 @@ def get_worker_report(self):
           "percent": self.process.get_cpu_percent(0)
         },
         "mem": {
-          "rss": self.process.get_memory_info().rss
+          "rss": self.get_memory()
         }
         # https://code.google.com/p/psutil/wiki/Documentation
         # get_open_files
@@ -233,7 +245,7 @@ def get_worker_report(self):
   def report_worker(self, w=0):
 
     try:
-      self.mongodb_logs.mrq_workers.update({
+      self.mongodb_jobs.mrq_workers.update({
         "_id": ObjectId(self.id)
       }, {"$set": self.get_worker_report()}, upsert=True, w=w)
     except pymongo.errors.AutoReconnect:
@@ -321,7 +333,14 @@ def work_loop(self):
       while True:
 
         while True:
+
+          if self.config["trace_memory"]:
+            # When debugging memory, intermediate psutils call like this one are
+            # needed for some obscure reason. (tested in test_memoryleaks.py)
+            self.get_memory()
+
           free_pool_slots = self.gevent_pool.free_count()
+
           if free_pool_slots > 0:
             self.status = "wait"
             break
@@ -387,6 +406,9 @@ def perform_job(self, job):
         This is the first call happening inside the greenlet.
     """
 
+    if self.config["trace_memory"]:
+      job.trace_memory_start()
+
     set_current_job(job)
 
     gevent_timeout = gevent.Timeout(job.timeout, JobTimeoutException(
@@ -437,6 +459,9 @@ def perform_job(self, job):
 
       self.done_jobs += 1
 
+      if self.config["trace_memory"]:
+        job.trace_memory_stop()
+
   def shutdown_graceful(self):
     """ Graceful shutdown: waits for all the jobs to finish. """
 

tests/conftest.py

@@ -9,6 +9,7 @@
 import time
 import re
 import json
+import urllib2
 
 sys.path.append(os.getcwd())
 
@@ -205,6 +206,13 @@ def send_task_cli(self, path, params, block=True, queue=None, **kwargs):
       return json.loads(out)
     return out
 
+  def get_report(self):
+    wait_for_net_service("127.0.0.1", 20020)
+    f = urllib2.urlopen("http://127.0.0.1:20020")
+    data = json.load(f)
+    f.close()
+    return data
+
 
 class RedisFixture(ProcessFixture):
   def flush(self):
@@ -214,9 +222,10 @@ def flush(self):
 class MongoFixture(ProcessFixture):
   def flush(self):
     for mongodb in (connections.mongodb_jobs, connections.mongodb_logs):
-      for c in mongodb.collection_names():
-        if not c.startswith("system."):
-          mongodb.drop_collection(c)
+      if mongodb:
+        for c in mongodb.collection_names():
+          if not c.startswith("system."):
+            mongodb.drop_collection(c)
 
 
 @pytest.fixture(scope="function")

tests/test_general.py

@@ -12,7 +12,7 @@ def test_general_simple_task_one(worker):
 
   time.sleep(0.1)
 
-  db_workers = list(worker.mongodb_logs.mrq_workers.find())
+  db_workers = list(worker.mongodb_jobs.mrq_workers.find())
   assert len(db_workers) == 1
   assert db_workers[0]["status"] in ["full", "wait"]
 
@@ -40,7 +40,7 @@ def test_general_simple_task_one(worker):
   assert db_jobs[0]["time"] < 0.1
   assert db_jobs[0]["switches"] >= 1
 
-  db_workers = list(worker.mongodb_logs.mrq_workers.find())
+  db_workers = list(worker.mongodb_jobs.mrq_workers.find())
   assert len(db_workers) == 1
   assert db_workers[0]["_id"] == db_jobs[0]["worker"]
   assert db_workers[0]["status"] == "stop"