django-q2/django_q/q.py

import importlib
from random import randint
from multiprocessing import Queue, Process, Event, current_process, cpu_count
import sys
import signal
import logging
from time import sleep

import jsonpickle as json

import coloredlogs
import redis
from django.conf import settings
from django.utils import timezone

from .models import Task

from .humanhash import uuid

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

logger = logging.getLogger('django-q')
coloredlogs.install(level=logging.INFO)


def test():
    for i in range(20):
        defer(u'testq.tasks.multiply', 5, i * randint(2, 100))


def defer(func, *args, **kwargs):
    # [name, func, args, kwargs, started, finished, result]
    pack = json.dumps([uuid()[0], func, args, kwargs, timezone.now()])
    r.rpush(q_list, pack)
    logger.debug('Pushed {}'.format(pack))


class Worker(object):
    def __init__(self):
        signal.signal(signal.SIGTERM, self.sig_handler)
        signal.signal(signal.SIGINT, self.sig_handler)
        self.running = True
        self.stable_size = cpu_count()
        self.stable = []
        self.task_queue = Queue()
        self.done_queue = Queue()
        self.fail_queue = Queue()
        # Spawn work horses
        for i in range(self.stable_size):
            self.spawn_horse()
        # Spawn monitors
        self.success_monitor_pid = None
        self.spawn_success_monitor()
        self.failure_monitor_pid = None
        self.spawn_failure_monitor()
        # Spawn pusher
        self.pusher_pid = None
        self.pusher_stop = Event()
        self.spawn_pusher()
        # Monitor process health
        while self.running:
            self.stable_boy()
            sleep(1)

    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_pusher(self):
        self.pusher_pid = self.spawn_process(self.pusher, self.task_queue, self.pusher_stop)

    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()
            logger.warn("reincarnated success monitor after death of {}".format(pid))
        elif pid == self.failure_monitor_pid:
            self.spawn_failure_monitor()
            logger.warn("reincarnated failure monitor after death of {}".format(pid))
        else:
            self.spawn_horse()
            logger.warn("reincarnated work horse after death of {}".format(pid))

    @staticmethod
    def pusher(task_queue, e):
        logger.info('{} pushing tasks at {}'.format(current_process().name, current_process().pid))
        while not e.is_set():
            task = r.blpop(q_list, 1)
            if task:
                task_queue.put(task[1])
                logger.debug('queueing {}'.format(task[1]))

    @staticmethod
    def success_monitor(done_queue):
        name = current_process().name
        logger.info("{} monitoring successes at {}".format(name, current_process().pid))
        for task in iter(done_queue.get, 'STOP'):
            logger.info("Finished [{}:{}]".format(task[1], task[0]))
            Task.objects.create(name=task[0], func=task[1], task=json.dumps(task),
                                started=task[4],
                                stopped=task[5])
        logger.info("{} stopped".format(name))

    @staticmethod
    def failure_monitor(fail_queue):
        name = current_process().name
        logger.info("{} monitoring failures at {}".format(name, current_process().pid))
        for task in iter(fail_queue.get, 'STOP'):
            logger.error("Failure [{}:{} - {}]".format(task[1], task[0], task[6]))
            Task.objects.create(name=task[0], func=task[1], task=json.dumps(task),
                                started=task[4],
                                stopped=task[5], success=False)
        logger.info("{} stopped".format(name))

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

    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
        logger.info('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')
            elif p.pid == self.pusher_pid:
                self.pusher_stop.set()
            else:
                self.task_queue.put('STOP')
            p.join()

        logger.info('Goodbye. Have a wonderful time.')

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