django-q2/django_q/apps.py

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

from django.apps import AppConfig
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)


class SessionAdminConfig(AppConfig):
    name = 'django_q'
    verbose_name = "Django Q"


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


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()
        # Spawn work horses
        for i in range(self.stable_size):
            self.spawn_horse()
        # Spawn monitor
        self.monitor_pid = None
        self.spawn_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)

    def spawn_monitor(self):
        self.monitor_pid = self.spawn_process(self.monitor, self.done_queue)

    def reincarnate(self, pid):
        if pid == self.monitor_pid:
            self.spawn_monitor()
            logger.warn("reincarnated monitor after death of {}".format(pid))
        elif pid == self.pusher_pid:
            self.spawn_pusher()
            logger.warn("reincarnated pusher 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 monitor(done_queue):
        name = current_process().name
        logger.info("{} monitoring at {}".format(name, current_process().pid))
        for task in iter(done_queue.get, 'STOP'):
            name = task[0]
            func = task[1]
            result = task[6]
            success = task[7]
            if success:
                logger.info("Finished [{}:{}]".format(func, name))
                result = json.dumps(result)
            else:
                logger.error("Failed [{}:{}] - {}".format(func, name, result))
            Task.objects.create(name=name,
                                func=func,
                                task=json.dumps(task),
                                started=task[4],
                                stopped=task[5],
                                result=result,
                                success=success)
        logger.info("{} stopped".format(name))

    @staticmethod
    def horse(task_queue, done_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)
                task.append(True)
                done_queue.put(task)
            except Exception as e:
                task.append(e)
                task.append(False)
                done_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.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()