subprocess_manager.py

import gevent
from gevent.event import Event
from gevent import subprocess
from gevent import signal as gevent_signal

import os
import sys
import signal



class SubProcess(object):
    def __init__(self, *args, **kwargs):
        env = {'WORKER': '1'}
        env.update(kwargs.get('env', {}))
        kwargs['env'] = env

        self.process = subprocess.Popen(*args, **kwargs)
        print "SubProcess start"

    def __getattr__(self, name):
        return getattr(self.process, name)


class SubProcessManager(object):
    def __init__(self, evt):
        self.processes = {}
        self.evt = evt
        # gevent.spawn(self._auto_restart)
        self.start_check_env = Event()
        gevent.spawn(self._check)

    def close(self):
        self.start_check_env.set()

    def start_subprocess(self, *args, **kwargs):
        p = SubProcess(*args, **kwargs)
        self.processes[p] = {'args': args, 'kwargs': kwargs}

    def kill_all(self):
        for p in self.processes.keys():
            p.kill()

    def restart_all(self):
        self.kill_all()
        params = self.processes.values()
        self.processes = {}

        for param in params:
            self.start_subprocess(*param['args'], **param['kwargs'])

    def _auto_restart(self):
        while True:
            for p, param in self.subprocesses.items():
                if p.poll():
                    print "{0} has done, code = {1}".format(p.pid, p.returncode)
                    
            gevent.sleep(1)


    def _check(self):
        self.start_check_env.wait()
        while not all([p.poll() is not None for p in self.processes.keys()]):
            gevent.sleep(0.1)

        print "check finish, exit"
        self.evt.set()

    

class Process(object):
    def __init__(self, master_env_key='WORKER'):
        self.is_master = os.environ.get(master_env_key, '0') == '0'
        if self.is_master:
            self.name = 'MASTER'
            self.evt = Event()
            self.manager = SubProcessManager(self.evt)
        else:
            self.name = 'SLAVE'

        self._set_signal_handler()
        print "start {0} process".format(self.name), os.getpid(), os.getppid()

    def _set_signal_handler(self):
        if self.is_master:
            def restart_handler(*args):
                print "restart"
                self.manager.restart_all()
                print "restart done"

            def quit_handler(*args):
                print "quit"
                self.quit(exit_code=1, kill=True)

            gevent_signal(signal.SIGHUP, restart_handler)
            gevent_signal(signal.SIGQUIT, quit_handler)

        else:
            pass

    def quit(self, exit_code=0, kill=False):
        if not self.is_master:
            return

        print "quit", exit_code
        if kill:
            self.manager.kill_all()
        sys.exit(exit_code)


    def run(self):
        if self.is_master:
            command = [sys.executable, __file__]

            self.manager.start_subprocess(command)
            self.manager.close()

            def wait_for_exit():
                self.evt.wait()
                self.quit()

            gevent.spawn(wait_for_exit)
            gevent.wait()
        else:
            # work
            for i in range(20):
                print i
                gevent.sleep(1)


if __name__ == '__main__':
    p = Process()
    p.run()