new_workers.py

from gevent import monkey
monkey.patch_socket()
import gevent.pool

import os
import signal
import random
import time
import datetime
from multiprocessing import Semaphore, Array



_signames = dict((getattr(signal, signame), signame) \
                    for signame in dir(signal) \
                    if signame.startswith('SIG') and '_' not in signame)


def signal_name(signum):
    # Hackety-hack-hack: is there really no better way to reverse lookup the
    # signal name?  If you read this and know a way: please provide a patch :)
    try:
        return _signames[signum]
    except KeyError:
        return 'SIG_UNKNOWN'


class BaseWorker(object):

    def __init__(self, num_processes=1):
        self._is_stopped = False

    def work(self):
        self._install_signal_handlers()
        while True:
            if self._is_stopped:
                break            
            self.spawn_child()
        self.quit()

    def quit(self):
        """ Enters a loop to wait for all children to finish and then quit """
        while True:
            if not self.has_active_horses():
                break
            time.sleep(1)

    def spawn_child(self):
        raise NotImplementedError('Implement this in a subclass.')

    def has_active_horses(self):
        # Each worker class has to implement a way of checking whether it is
        # in the middle of running one or more jobs
        raise NotImplementedError('Implement this in a subclass.') 

    def handle_quit(self):
        print 'QUITTT'

    def fake_work(self):        
        # When working, children should ignore CTRL+C
        signal.signal(signal.SIGINT, signal.SIG_IGN)
        signal.signal(signal.SIGTERM, signal.SIG_DFL)

        sleep_time = 5 * random.random()
        print datetime.datetime.now(), '- Hello from', os.getpid(), '- %.3fs' % sleep_time
        time.sleep(sleep_time)
        print datetime.datetime.now(), '- Done from', os.getpid()

    def _install_signal_handlers(self):
        """Installs signal handlers for handling SIGINT and SIGTERM
        gracefully.
        """

        def request_force_stop(signum, frame):
            """Terminates the application (cold shutdown).
            """
            print 'Cold shut down.'
            # Need to call ``handle_cold_shutdown`` implemented by subclasses
            #self.handle_cold_shutdown()
            raise SystemExit()

        def request_stop(signum, frame):
            """Stops the current worker loop but waits for child processes to
            end gracefully (warm shutdown).
            """
            self._is_stopped = True
            print '%s Got signal %s.' % (os.getpid(), signal_name(signum))
            signal.signal(signal.SIGINT, request_force_stop)
            signal.signal(signal.SIGTERM, request_force_stop)
            print 'Warm shut down requested.'
            
        signal.signal(signal.SIGINT, request_stop)
        signal.signal(signal.SIGTERM, request_stop)    


def process_is_alive(pid):
    # Check if a process is alive by sending it a signal that does nothing
    # If OSError is raised, it means the process is no longer alive
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False


class ForkingWorker(BaseWorker):

    def __init__(self, num_processes=1):
        super(ForkingWorker, self).__init__()
        # Set up sync primitives, to communicate with the spawned children
        self._semaphore = Semaphore(num_processes)
        self._slots = Array('i', [0] * num_processes)

    def spawn_child(self):
        """Forks and executes the job."""
        
        self._semaphore.acquire()    # responsible for the blocking
        # If CTRL+C is pressed when blocking to acquire children
        # (causes self._stopped to be set to True) return right away
        if self._is_stopped:
            return


        # Select an empty slot from self._slots (the first 0 value is picked)
        # The implementation guarantees there will always be at least one empty slot
        for slot, value in enumerate(self._slots):
            if value == 0:
                break

        # The usual hardcore forking action
        child_pid = os.fork()
        if child_pid == 0:
            random.seed()
            # Within child
            try:
                self.fake_work()
            finally:
                # This is the new stuff.  Remember, we're in the child process
                # currently. When all work is done here, free up the current
                # slot (by writing a 0 in the slot position).  This
                # communicates to the parent that the current child has died
                # (so can safely be forgotten about).
                self._slots[slot] = 0
                self._semaphore.release()
                os._exit(0)
        else:
            # Within parent, keep track of the new child by writing its PID
            # into the first free slot index.
            self._slots[slot] = child_pid

    def has_active_horses(self):
        # If any of the worker slot is non zero, that means there's a job still running
        for pid in self._slots:            
            if pid and process_is_alive(pid):
                return True
        return False