test-keep-alive-server-so

#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
#
# <header-blahblahblah>

import sys, functools, socket, logging

SERVER_FAMILY = socket.AF_INET # socket.AF_INET6 or socket.AF_INET  ?
SERVER_ADDR = '0.0.0.0'
SERVER_PORT = 8124
SERVER_SO_KEEPALIVE = True

from tornado import ioloop
from tornado import iostream
from tornado import netutil

def new_peer_name():
    from os import urandom
    from hashlib import sha256
    
    h = sha256()
    h.update(urandom(100))
    peer_name = 'peer_' + h.hexdigest()[0:10]
    
    return peer_name

class Peer(object):
    def __init__(self, peer_name, connection, address, peer_counter):
        self._log = logging.getLogger('test-alive-server-peer')
        self._peer_name = peer_name
        self._address = address
        self._peer_counter = peer_counter
        
        self._peer_counter.count += 1
        
        self._log.info('accepted! {}'.format(dict(
            peer_name=self._peer_name,
            address=self._address,
            all_peers_count=self._peer_counter.count,
            sockopt_keepalive=connection.getsockopt(
                    socket.SOL_SOCKET, socket.SO_KEEPALIVE),
        )))
        
        self._stream = iostream.IOStream(connection)
        self._stream.set_close_callback(self._close_handle)
        self._stream.read_until_close(
            functools.partial(self._read_handle, last=True),
            streaming_callback=self._read_handle,
        )
    
    def _close_handle(self):
        self._peer_counter.count -= 1
        
        self._log.info('closed! {}'.format(dict(
            peer_name=self._peer_name,
            address=self._address,
            all_peers_count=self._peer_counter.count,
        )))
    
    def _read_handle(self, data, last=None):
        if last is None:
            last = False
        
        self._log.info('read: {}'.format(dict(
            data=data,
            last=last,
            peer_name=self._peer_name,
            address=self._address,
        )))
        
        if not self._stream.closed():
            self._stream.write('echo {}\n'.format(dict(
            peer_name=self._peer_name,
            address=self._address,
            data=data,
        )))


class PeerCounter(object):
    def __init__(self):
        self.count = 0

class TestKeepAliveServer(object):
    def __init__(self):
        self._log = logging.getLogger('test-alive-server')
        self._log.info('begin test-keep-alive-server')
        
        self._peer_counter = PeerCounter()
        
        try:
            socks = netutil.bind_sockets(
                    SERVER_PORT, address=SERVER_ADDR, family=SERVER_FAMILY)
        except IOError:
            socks = None
        if not socks:
            self._log.info('can not bind')
            self._exit()
            return
        
        sock = socks[0]
        if SERVER_SO_KEEPALIVE:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        netutil.add_accept_handler(sock, self._accept_handler)
    
    def _exit(self):
        self._log.info('exit test-keep-alive-server')
        
        io_loop = ioloop.IOLoop.instance()
        io_loop.stop()
    
    def _accept_handler(self, connection, address):
        peer_name = new_peer_name()
        
        Peer(peer_name, connection, address, self._peer_counter)

def main():
    logging.basicConfig(
            filename='test-keep-alive-server.log',
            format='%(levelname)s %(asctime)-15s %(name)s: %(message)s',
            level=logging.DEBUG)
        
    TestKeepAliveServer()
    
    io_loop = ioloop.IOLoop.instance()
    io_loop.start()

if __name__ == '__main__':
    error = main()
    if error:
        exit(error)