django-q2/django_q/q.py

import importlib
from time import sleep
from multiprocessing import Queue, Process, current_process, cpu_count
import sys
import signal
from uuid import uuid4
import ujson as json

import redis
from django.conf import settings

r = redis.StrictRedis(decode_responses=True)
secret = settings.SECRET_KEY
prefix = 'django_q'
q_list = '{}:q'.format(prefix)


def test():
    for i in range(4):
        defer(u'testq.tasks.multiply', 5, i)


def defer(func, *args, **kwargs):
    pack = json.dumps([uuid4().urn, func, args, kwargs])
    r.rpush(q_list, pack)


class Worker(object):
    def __init__(self):
        self.running = True
        self.stable_size = cpu_count() - 2
        self.stable = []
        self.task_queue = Queue()
        self.done_queue = Queue()
        self.fail_queue = Queue()
        self.failure_monitor_pid = None
        self.success_monitor_pid = None
        # Spawn work horses
        for i in range(self.stable_size):
            self.spawn_horse()
        # Spawn monitors
        self.spawn_success_monitor()
        self.spawn_failure_monitor()
        # Attach signal handler
        signal.signal(signal.SIGTERM, self.sig_handler)
        signal.signal(signal.SIGINT, self.sig_handler)
        # Keep popping Redis
        while self.running:
            self.stable_boy()
            sleep(0.2)
            task = r.lpop(q_list)
            if task:
                self.task_queue.put(task)

    def spawn_process(self, target, *args):
        # This is just for PyCharm to not crash. Ignore it.
        if not hasattr(sys.stdin, 'close'):
            def dummy_close():
                pass

            sys.stdin.close = dummy_close
        p = Process(target=target, args=args)
        self.stable.append(p)
        p.start()
        return p.pid

    def spawn_horse(self):
        self.spawn_process(self.horse, self.task_queue, self.done_queue, self.fail_queue)

    def spawn_success_monitor(self):
        self.success_monitor_pid = self.spawn_process(self.success_monitor, self.done_queue)

    def spawn_failure_monitor(self):
        self.failure_monitor_pid = self.spawn_process(self.failure_monitor, self.fail_queue)

    def reincarnate(self, pid):
        if pid == self.success_monitor_pid:
            self.spawn_success_monitor()
        elif pid == self.failure_monitor_pid:
            self.spawn_failure_monitor()
        else:
            self.spawn_horse()

    @staticmethod
    def success_monitor(done_queue):
        name = current_process().name
        print("{} monitoring successes at {}".format(name, current_process().pid))
        for result in iter(done_queue.get, 'STOP'):
            task = result[0]
            res = result[1]
            print("Success [{}:{} - {}]".format(task[1], task[0], res))
        print("{} stopped".format(name))

    @staticmethod
    def failure_monitor(fail_queue):
        name = current_process().name
        print("{} monitoring failures at {}".format(name, current_process().pid))
        for result in iter(fail_queue.get, 'STOP'):
            task = result[0]
            e = result[1]
            print("Failure [{}:{} - {}]".format(task[1], task[0], e))
        print("{} stopped".format(name))

    @staticmethod
    def horse(task_queue, done_queue, fail_queue):
        name = current_process().name
        print('{} ready for work at {}'.format(name, current_process().pid))
        for pack in iter(task_queue.get, 'STOP'):
            task = json.loads(pack)
            uid = task[0]
            func = task[1]
            module, func = func.rsplit('.', 1)
            args = task[2]
            kwargs = task[3]
            print(name, 'Processing [{}:{}]'.format(func, uid))
            try:
                m = importlib.import_module(module)
                f = getattr(m, func)
                result = f(*args, **kwargs)
                done_queue.put((task, result))
            except TypeError as e:
                fail_queue.put((task, e))
        print(name, 'Stopped')

    def stable_boy(self):
        # Check if all the horses are alive
        for p in list(self.stable):
            if not p.is_alive():
                # Be humane
                p.terminate()
                self.stable.remove(p)
                # Replace it with a fresh one
                self.reincarnate(p.pid)

    def stop(self):
        # Send the STOP signal to the stable
        self.running = False
        print('Stopping')
        # Wait for all the workers to finish the queue
        for p in self.stable:
            if p.pid == self.failure_monitor_pid:
                self.fail_queue.put('STOP')
            elif p.pid == self.success_monitor_pid:
                self.done_queue.put('STOP')
            else:
                self.task_queue.put('STOP')
            p.join()

        print('Goodbye. Have a wonderful time.')

    def sig_handler(self, signum, frame):
        self.stop()