Passing the redis connection around the cluster to reduce connections

django_q/core.py

@@ -50,17 +50,11 @@ if not logger.handlers:
     handler.setFormatter(formatter)
     logger.addHandler(handler)
 
-# Redis
-r = redis.StrictRedis(**REDIS)
+redis_client = redis.StrictRedis(**REDIS)
 
 
 class Cluster(object):
     def __init__(self, list_key=Q_LIST):
-        try:
-            r.ping()
-        except ():
-            logger.error('Can not connect to Redis server')
-            return
         self.sentinel = None
         self.stop_event = None
         self.start_event = None
@@ -139,6 +133,7 @@ class Sentinel(object):
         self.parent_pid = os.getppid()
         self.name = current_process().name
         self.list_key = list_key
+        self.r = redis_client
         self.status = None
         self.reincarnations = 0
         self.tob = timezone.now()
@@ -172,7 +167,7 @@ class Sentinel(object):
         return p.pid
 
     def spawn_pusher(self):
-        return self.spawn_process(pusher, self.task_queue, self.event_out, self.list_key)
+        return self.spawn_process(pusher, self.task_queue, self.event_out, self.list_key, self.r)
 
     def spawn_worker(self):
         self.spawn_process(worker, self.task_queue, self.done_queue)
@@ -246,7 +241,7 @@ class Sentinel(object):
         Stat(self, message).save()
 
 
-def pusher(task_queue, e, list_key=Q_LIST):
+def pusher(task_queue, e, list_key=Q_LIST, r=None):
     """
     Pulls tasks of the Redis List and puts them in the task queue
     :type task_queue: multiprocessing.Queue
@@ -254,6 +249,8 @@ def pusher(task_queue, e, list_key=Q_LIST):
     :type list_key: str
     """
     logger.info('{} pushing tasks at {}'.format(current_process().name, current_process().pid))
+    if not r:
+        r = redis_client
     while True:
         task = r.blpop(list_key, 1)
         if task:
@@ -364,6 +361,11 @@ def async(func, *args, **kwargs):
         del kwargs['list_key']
     else:
         list_key = Q_LIST
+    if 'redis' in kwargs:
+        r = kwargs['redis']
+        del kwargs['redis']
+    else:
+        r = redis_client
     task = {'name': uuid()[0], 'func': func, 'hook': hook, 'args': args, 'kwargs': kwargs, 'started': timezone.now()}
     pack = SignedPackage.dumps(task)
     r.rpush(list_key, pack)
@@ -436,6 +438,7 @@ class Stat(Status):
         super(Stat, self).__init__(sentinel.parent_pid)
         if message:
             sentinel.status = message
+        self.r = sentinel.r
         self.tob = sentinel.tob
         self.reincarnations = sentinel.reincarnations
         self.sentinel = sentinel.pid
@@ -466,18 +469,20 @@ class Stat(Status):
         return '{}:cluster:{}'.format(PREFIX, cluster_id)
 
     def save(self):
-        r.set(self.key, SignedPackage.dumps(self, True), 3)
+        self.r.set(self.key, SignedPackage.dumps(self, True), 3)
 
     def empty_queues(self):
         return self.done_q_size + self.task_q_size == 0
 
     @staticmethod
-    def get(cluster_id):
+    def get(cluster_id, r=None):
         """
         gets the current status for the cluster
         :param cluster_id: id of the cluster
         :return: Stat or Status
         """
+        if not r:
+            r = redis_client
         key = Stat.get_key(cluster_id)
         if r.exists(key):
             pack = r.get(key)
@@ -488,11 +493,13 @@ class Stat(Status):
         return Status(cluster_id)
 
     @staticmethod
-    def get_all():
+    def get_all(r=None):
         """
         Gets status for all currently running clusters with the same prefix and secret key
         :return: Stat list
         """
+        if not r:
+            r = redis_client
         stats = []
         keys = r.keys(pattern='{}:cluster:*'.format(PREFIX))
         if keys:
@@ -504,6 +511,12 @@ class Stat(Status):
                     continue
         return stats
 
+    def __getstate__(self):
+        # Don't pickle the redis connection
+        state = dict(self.__dict__)
+        del state['r']
+        return state
+
 
 def scheduler():
     """

django_q/management/commands/qmonitor.py

@@ -6,7 +6,7 @@ from django.utils import timezone
 from blessed import Terminal
 
 # Local
-from django_q.core import Stat, RUNNING, STOPPED
+from django_q.core import Stat, RUNNING, STOPPED, redis_client
 
 
 class Command(BaseCommand):
@@ -15,8 +15,10 @@ class Command(BaseCommand):
     def handle(self, *args, **options):
         monitor()
 
+
 def monitor(run_once=False):
     term = Terminal()
+    r = redis_client
     with term.fullscreen(), term.hidden_cursor(), term.cbreak():
         val = None
         start_width = int(term.width / 8)
@@ -35,7 +37,7 @@ def monitor(run_once=False):
             print(term.move(0, 6 * col_width) + term.black_on_green(term.center('Deaths', width=col_width - 1)))
             print(term.move(0, 7 * col_width) + term.black_on_green(term.center('Uptime', width=col_width - 1)))
             i = 2
-            stats = Stat.get_all()
+            stats = Stat.get_all(r=r)
             print(term.clear_eos())
             for stat in stats:
                 # color status
@@ -63,6 +65,6 @@ def monitor(run_once=False):
                 i += 1
             # for testing
             if run_once:
-                return Stat.get_all()
+                return Stat.get_all(r=r)
             print(term.move(i + 2, 0) + term.center('[Press q to quit]'))
             val = term.inkey(timeout=1)

django_q/tests/test_cluster.py

@@ -3,11 +3,12 @@ import os
 from multiprocessing import Queue, Event
 
 import pytest
+from conf import REDIS
 
 myPath = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, myPath + '/../')
 
-from django_q.core import Cluster, r, async, pusher, worker, monitor, Sentinel
+from django_q.core import Cluster, async, pusher, worker, monitor, Sentinel
 from django_q.humanhash import DEFAULT_WORDLIST
 from django_q import result, get_task, Task
 from django_q.tests.tasks import multiply
@@ -20,8 +21,12 @@ class WordClass(object):
     def get_words(self):
         return self.word_list
 
+@pytest.fixture
+def r():
+    import redis
+    return redis.StrictRedis(**REDIS)
 
-def test_redis_connection():
+def test_redis_connection(r):
     assert r.ping() is True
 
 
@@ -46,7 +51,7 @@ def test_sentinel():
 
 
 @pytest.mark.django_db
-def test_cluster():
+def test_cluster(r):
     list_key = 'cluster_test:q'
     r.delete(list_key)
     task = async('django_q.tests.tasks.count_letters', DEFAULT_WORDLIST, list_key=list_key)
@@ -58,7 +63,7 @@ def test_cluster():
     event = Event()
     event.set()
     # Test push
-    pusher(task_queue, event, list_key=list_key)
+    pusher(task_queue, event, list_key=list_key, r=r)
     assert task_queue.qsize() == 1
     assert r.llen(list_key) == 0
     # Test work
@@ -76,7 +81,7 @@ def test_cluster():
 
 
 @pytest.mark.django_db
-def test_async():
+def test_async(r):
     list_key = 'cluster_test:q'
     r.delete(list_key)
     a = async('django_q.tests.tasks.count_letters', DEFAULT_WORDLIST, hook='django_q.tests.test_cluster.assert_result',
@@ -111,7 +116,7 @@ def test_async():
     stop_event.set()
     # push the tasks
     for i in range(task_count):
-        pusher(task_queue, stop_event, list_key=list_key)
+        pusher(task_queue, stop_event, list_key=list_key, r=r)
     assert r.llen(list_key) == 0
     assert task_queue.qsize() == task_count
     task_queue.put('STOP')