twisted_respawning.py

import os
import os.path
import sys
import signal

from twisted.application import service
from twisted.internet import protocol, reactor
from twisted.python import log

from twisted.scripts._twistd_unix import ServerOptions

RESPAWN_SIGNAL = signal.SIGHUP

class ChildProtocol(protocol.ProcessProtocol):
    def connectionMade(self):
        log.msg('New service launched, PID %d' % self.transport.pid)
        self.transport.writeToChild(4, 'Data from old to new process')

    def childDataReceived(self, childFD, data):
        log.msg('[%d] %s' % (childFD, data), system='child')

    def processEnded(self, reason):
        log.msg('Process ended: %s' % reason, system='child')

def checkSignal(sig):
    if signal.getsignal(sig) != signal.SIG_DFL:
        raise RuntimeError('Signal handler already registered')

def registerSignalHandler(sig, fun):
    log.msg('Registering signal handler')
    signal.signal(sig, lambda *_: reactor.callLater(0.01, fun))

def unregisterSignalHandler(sig):
    log.msg('Unregistering signal handler')
    signal.signal(sig, signal.SIG_DFL)

def calculateGenerationPath(generation, original):
    if '.' in original:
        prefix, suffix = original.rsplit('.', 1)
        if suffix.isdigit():
            return '%s.%d' % (prefix, generation)
        else:
            return '%s.%d' % (original, generation)
    else:
        return '%s.%d' % (original, generation)


def spawnChild():
    generation = int(os.environ.get('AGENT_PROCESS_GENERATION', 0)) + 1

    proto = ChildProtocol()
    executable = sys.executable

    args = [executable, sys.argv[0]]

    # Don't pass pidfile or logfile
    filters = set()
    filters.add(lambda opt: not opt.startswith('--pidfile='))
    filters.add(lambda opt: not opt.startswith('--logfile='))
    apply_filters = lambda arg: all(map(lambda fun: fun(arg), filters))

    twistd_args = filter(apply_filters, sys.argv[1:])

    parser = ServerOptions()
    parser.parseOptions()

    if not 'pidfile' in parser or not parser['pidfile']:
        raise RuntimeError('Unable to locate pidfile')
    pidfile = calculateGenerationPath(generation, parser['pidfile'])

    if 'logfile' in parser and parser['logfile']:
        logfile = parser['logfile']
    else:
        logfile = 'respawn.log'
        
    logfile = calculateGenerationPath(generation, logfile)
    if os.path.exists(logfile):
        os.unlink(logfile)

    
    args.append('--pidfile=%s' % pidfile)
    args.append('--logfile=%s' % logfile)
    args.extend(twistd_args)
    args = tuple(args)

    env = os.environ.copy()
    env['AGENT_PROCESS_GENERATION'] = str(generation)

    proc = reactor.spawnProcess(proto, executable, args, childFDs={
        1: 'r',
        2: 'r',
        3: 'r',
        4: 'w',
    }, env=env)

def respawn(sig):
    log.msg('Respawning')
    unregisterSignalHandler(sig)

    spawnChild()

    reactor.callLater(2, reactor.stop)

class RespawnService(service.Service):
    def startService(self):
        log.msg('Starting respawn service')

        checkSignal(RESPAWN_SIGNAL)
        registerSignalHandler(RESPAWN_SIGNAL, lambda: respawn(RESPAWN_SIGNAL))

        # This is a demo...
        reactor.callLater(1, lambda: os.kill(os.getpid(), RESPAWN_SIGNAL))

        if 'AGENT_PROCESS_GENERATION' in os.environ:
            fd = os.fdopen(3, 'wb')
            fd.write('Sending from new to old process')
            fd.flush()

        return service.Service.startService(self)

    def stopService(self):
        log.msg('Stopping respawn service')
        unregisterSignalHandler(RESPAWN_SIGNAL)
        return service.Service.stopService(self)