Simple UDP load balancer

udplb.py

#!/usr/bin/env python
#------------------------------------------------------------------------------
# UDP Load Balancer featuring simple round robin with session affinity.
#------------------------------------------------------------------------------

import sys
import signal
import logging
from socket import *

LB_NODES = ['node1.example.com', '192.168.2.1']
LB_PORT = 2100

affinityTable = {}
_next = 0


def get_next_node(srcIP):
    global affinityTable, _next
    if srcIP not in affinityTable:
        _next = (_next + 1) % len(LB_NODES)
        affinityTable[srcIP] = _next
    return LB_NODES[affinityTable[srcIP]]


def main():
    def signal_handler(signal, frame):
        logging.info('Interrupted, shutting down')
        s1.close()
        s2.close()

    signal.signal(signal.SIGINT, signal_handler)

    logging.basicConfig(format='%(asctime)s [%(levelname)s] - %(message)s',
                        datefmt='%m/%d/%Y %I:%M:%S %p',
                        filename='udplb.log', level=logging.INFO)

    logging.info('Starting server on port %d' % LB_PORT)
    try:
        s1 = socket(AF_INET, SOCK_DGRAM)
        s1.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        s1.bind(('', LB_PORT))
    except:
        logging.exception('')
        sys.exit(1)

    bad_nodes = []
    for i, host in enumerate(LB_NODES):
        try:
            ip = gethostbyname(host)
            if LB_NODES[i] != ip:
                logging.info('Node %s resolved to IP %s' % (host, ip))
                LB_NODES[i] = ip

        except gaierror:
            logging.error('Name resolution failed for node %s, skipping' % host)
            bad_nodes.append(i)

    for i in bad_nodes:
        del(LB_NODES[i])

    s2 = socket(AF_INET, SOCK_DGRAM)
    while True:
        try:
            buf, (srcip, sport) = s1.recvfrom(65535)
            s2.sendto(buf, (get_next_node(srcip), LB_PORT))
        except:
            logging.exception('')

if __name__ == '__main__':
    main()