daemonized

#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
@date: 2012-11-09
@author: shell.xu
'''
import os, sys, time, signal, getopt, logging
from os import path
from contextlib import contextmanager

def initlog(lv, logfile=None):
    rootlog = logging.getLogger()
    if logfile: handler = logging.FileHandler(logfile)
    else: handler = logging.StreamHandler()
    handler.setFormatter(
        logging.Formatter(
            '%(asctime)s,%(msecs)03d (%(process)d)%(name)s[%(levelname)s]: %(message)s',
            '%H:%M:%S'))
    rootlog.addHandler(handler)
    rootlog.setLevel(lv)
    return rootlog

def daemonized(parentalive=False, pidfile=None, logfile=None, chdir=True):
    if pidfile and path.exists(pidfile):
        with open(pidfile) as fi: pid = int(fi.read().strip())
        if path.exists('/proc/%d/exe' % pid): return -1
    try:
        pid = os.fork()
        if pid > 0:
            if parentalive: return pid
            else: os._exit(0)
    except OSError: os._exit(1)
    os.setsid()
    os.umask(0)
    with open('/dev/null', 'r') as dev_null:
        os.dup2(dev_null.fileno(), 0)
        if logfile:
            with open(logfile, 'a') as f:
                os.dup2(f.fileno(), 1)
                os.dup2(f.fileno(), 2)
        else:
            os.dup2(dev_null.fileno(), 1)
            os.dup2(dev_null.fileno(), 2)
    if chdir:
        if not isinstance(chdir, basestring): chdir = '/'
        os.chdir(chdir)
    try:
        if os.fork() > 0: os._exit(0)
    except OSError: sys.exit(1)
    if pidfile:
        with open(pidfile, 'w') as fo: fo.write(str(os.getpid()))

@contextmanager
def lockfile(filename, share=False):
    logging.debug('locking %s' % filename)
    f = open(filename, 'r')
    fcntl.flock(f.fileno(), fcntl.LOCK_SH if share else fcntl.LOCK_EX)
    try: yield f
    finally:
        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
        f.close()
        logging.debug('unlock %s' % filename)

def kill_stand(pids, timeout=5):
    if not pids: return
    t_start = time.time()
    logging.debug('try term %s.' % str(pids))
    for pid in pids:
        try: os.kill(pid, signal.SIGTERM)
        except OSError: pass
    while (time.time() - t_start) < timeout and pids:
        try: pid, st = os.waitpid(-1, os.WNOHANG)
        except OSError: return
        if pid: pids.remove(pid)
        time.sleep(0.1)
    logging.debug('try kill %s.' % str(pids))
    for pid in pids:
        try: os.kill(pid, signal.SIGKILL)
        except OSError: pass

def runner_wrap(runner):
    def inner(origpid):
        pid = os.fork()
        if pid > 0: return pid
        signal.signal(signal.SIGTERM, signal.SIG_DFL)
        st = 0
        try: st = runner(origpid)
        finally: os._exit(st)
    return inner

class SIGHUPError(StandardError): pass

class Watcher(object):

    def __init__(self, runners, timeout=2, lifetime=None):
        self.timeout, self.times, self.idlefunc, self.pids = timeout, 0, [], {}
        for r in runners: self.register(r)
        if lifetime: self.regidle(lifetime)(self.killall)

    def register(self, func): self.pids[func] = 0

    def run(self):
        signal.signal(signal.SIGTERM, self.sig_handler)
        signal.signal(signal.SIGHUP, self.sig_handler)
        for r, p in self.pids.iteritems(): self.pids[r] = r(p)
        try:
            while self.pids:
                try: self.waitpid()
                except KeyboardInterrupt: break
                except SIGHUPError: break
                except: logging.exception('watcher unknown exception')
        finally:
            logging.info('system exit')
            self.killall()

    def waitpid(self):
        time.sleep(self.timeout)
        try: pid, st = os.waitpid(-1, os.WNOHANG)
        except OSError: return
        if not pid: self.on_idle()
        else:
            for r, p in self.pids.iteritems():
                if p == pid:
                    self.pids[r] = r(pid)
                    break

    def killall(self): kill_stand(self.pids.values())

    def sig_handler(self, signum, frame):
        if signum == signal.SIGTERM:
            logging.info('signal TERM, begin to stop childs')
            tokill = self.pids.values()
            self.pids.clear()
            kill_stand(tokill)
        elif signum == signal.SIGHUP:
            logging.info('signal HUP, restart childs')
            raise SIGHUPError()

    # FIXME: sometimes, child die without a zombie
    def regidle(self, times=30):
        def recv(func):
            self.idlefunc.append((times, func))
            return func
        return recv

    def on_idle(self):
        self.times += 1
        for t, f in self.idlefunc:
            if self.times % t == 0: f()
        if self.times >= 0x100000000: self.times = 1

def Runner(args):
    def run(pre_pid):
        if pre_pid: logging.warn('prior proc stopped, pid %d' % pre_pid)
        pid = os.spawnvp(os.P_NOWAIT, args[0], args)
        logging.info('proc started with pid %d' % pid)
        return pid
    return run

def main():
    '''
    -f: foreground
    -l: log file
    -h: help
    -p: pidfile
    '''
    optlist, args = getopt.getopt(sys.argv[1:], 'fl:hp:')
    optdict = dict(optlist)
    if '-h' in optdict:
        print main.__doc__
        return

    if '-f' not in optdict:
        name = path.basename(args[0])
        pidfile = path.realpath(
            path.expanduser(optdict.get('-p', '/var/run/%s.pid' % name)))
        logfile = path.realpath(
            path.expanduser(optdict.get('-l', '/var/log/%s.log' % name)))
        daemonized(pidfile=pidfile, logfile=logfile)
    initlog(logging.WARNING)
    logging.info('daemoniz finished, new pid %d' % os.getpid())

    w = Watcher([Runner(args),])
    w.run()

if __name__ == '__main__':
    while True:
        try:
            main()
            break
        except SIGHUPError: pass