diff --git a/django_q/core.py b/django_q/core.py index a1e7a2f..f383e45 100644 --- a/django_q/core.py +++ b/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(): """ diff --git a/django_q/management/commands/qmonitor.py b/django_q/management/commands/qmonitor.py index 2673746..153f9b6 100644 --- a/django_q/management/commands/qmonitor.py +++ b/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) diff --git a/django_q/tests/test_cluster.py b/django_q/tests/test_cluster.py index 5e893f5..2bcb66a 100644 --- a/django_q/tests/test_cluster.py +++ b/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')