testclient.py

"""
You may wonder what the heck is going on here.
This is simply my take on code reloading with persisting open sockets i thought
of using in some future IRC bot. A good way to research a little bit more about
exec, file descriptors and sockets.
"""

import socket
import random
import time
import sys
import os
import atexit
import stat
import signal

instance_id = 'worker-%d' % random.randint(100, 999)
def l(m): sys.stderr.write('[%d/%s] %s\n' % (os.getpid(), instance_id, m))

l('starting')

def restart_process():
    # We don't really have to pass environment explicitly there since we store
    # the fd in os.environ already

    # Oh, oslo, this restarter is rather stupid, because it just throws a file at
    # an interpreter, without even checking if it exists (TODO: nginx-style?)
    os.execvp('python', ['python'] + sys.argv)

do_restart = True

def atexit_hook():
    if do_restart:
        l("atexit hook (1sec throttle)")
        time.sleep(1)
        restart_process()

atexit.register(atexit_hook)

def signal_handler(signal, frame):
    l('signal received: %d' % signal)
    restart_process()

s = None

if 'SOCKET_FD' in os.environ:
    try:
        l('fromfd(%s)' % os.environ['SOCKET_FD'])
        fd = int(os.environ['SOCKET_FD'])
        if stat.S_ISSOCK(os.fstat(fd).st_mode): # apparently, fromfd is *that* dumb
            s = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
            # This is the place where we should store new fileid into the env,
            # in order to atexit work properly on crashes between closing old
            # updating stored one

            os.environ['SOCKET_FD'] = str(s.fileno())
            # In worst case, if code fails exactly here we will just get one
            # unbound descriptor
            os.close(fd) # fromfd does dup, so we may want to close the original one
        else:
            l("it's not a socket, idiot")

    except:
        l('well, fromfd failed (%r)' % sys.exc_info()[1])
        s = None

if not s:
    l('new socket')
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(("localhost", 9999))
    os.environ['SOCKET_FD'] = str(s.fileno())

signal.signal(signal.SIGUSR1, signal_handler)

try:
    for i in range(500):
        s.send("%s: %d" % (instance_id, time.time()))
        sys.stderr.write('.'); sys.stderr.flush()
        time.sleep(0.5)

    l('bang!')
    restart_process()
except socket.error as e:
    if e.errno == 32:
        l('broken pipe, clear and restart')
        s.close()
        del os.environ['SOCKET_FD']
        restart_process()

except KeyboardInterrupt:
    l('exiting')
    s.close()
    do_restart = False

testserver.py

import SocketServer
import time

class MyTCPHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        s_time = time.time()
        self.log('connected')

        while True:
            data = self.request.recv(1024)
            if not data: break
            self.log(repr(data))

        self.log('disconnected after %.2f seconds' % (time.time() - s_time))

    def log(self, msg):
        print '[%s:%d] %s' % (self.client_address[0], self.client_address[1], msg)

if __name__ == "__main__":
    HOST, PORT = "localhost", 9999

    server = SocketServer.TCPServer((HOST, PORT), MyTCPHandler)
    server.serve_forever()