lem_server.py

Go to the documentation of this file.

#!/usr/bin/env python

# Here's some useful documentation on TCP: http://linux.die.net/man/7/tcp
# The python sockets howto is good https://docs.python.org/2/howto/sockets.html

# We need more open files than is usual. Edit /etc/security/limits.conf to add
# This is a per-login setting, no need to reboot.
# lem soft nofile 4096
# lem hard nofile 4096

# We need bigger tables to keep track of old connections. Absence of this is
# what made the network/CPU "oscillate" at high load. Add these commands to
# /etc/rc.d/rc.local
#
# touch /var/lock/subsys/local
# sysctl -w net.netfilter.nf_conntrack_max=262144
# echo 65536 > /sys/module/nf_conntrack/parameters/hashsize
#
# Check they took with:
#
# $ cat /sys/module/nf_conntrack/parameters/hashsize
# 65536
# $ cat /proc/sys/net/netfilter/nf_conntrack_max
# 262144

# If we ever want to listen on port 80 it'd be something like:
# ls -l /etc/authbind/byport/80
# -rwxr--r-- 1 lem     lem    0 Nov 18 20:36 /etc/authbind/byport/80
# and start with 'authbind --deep ./lem_server.py'

import os
import select
import socket
import sys
import time
import thread
import time
import traceback

from collections import defaultdict, deque

gVerbosity = 0

# IP whitelist is currently unused
gIPWhitelist = [
    # Fermilab, according to https://ipinfo.io/AS3152
    '131.225.',
    '192.190.216.',
    '198.49.208.',
    # Caltech, according to https://ipinfo.io/AS31
    # TODO maybe only the HEP cluster machines?
    '131.215.',
    '134.4.',
    '192.12.19.',
    '192.31.43.',
    '192.31.208.',
    '192.43.243.',
    '192.54.249.'
    # SMU: https://ipinfo.io/AS1832
    # UMN: https://ipinfo.io/AS217
    ]

def is_allowed_ip(ip):
    # We'll be running offsite. Don't know all those IP addresses, so just let
    # anything through for now.
    return True

    for entry in gIPWhitelist:
        if ip.startswith(entry): return True
    return False


class ConnInfo:
    def __init__(self, client, query, timeout):
        self.client = client
        self.query = query
        self.deadline = time.time()+int(timeout)

    def late(self):
        # Discretize so that we close several at second-boundaries
        # instead of one constantly.
        return int(time.time()) > int(self.deadline)



def pop_old_conns(q):
    ret = []
    while len(q) > 0 and q[0].late():
        ret.append(q.popleft())
    return ret

# Sums the length of all the lists in a dict
def nConns(d):
    return sum([len(x) for x in d.values()])

# Takes a dict of deques, removes the first element from every deque, and
# returns that list.
def pop_fronts(d):
    return [x.popleft() for x in d.values() if len(x) > 0]


# When a new 'pid' request comes in, if we can't match it to a worker
# immediately, we put it here. The top level key is the release, and under each
# release is a list of requests (ConnInfo objects), ordered by age, newest
# last.
gPIDQueue = defaultdict(deque)

# When a new 'get_work' request comes in, if we can't match it to a client
# immediately, we put it here. The top level key is the release, and under each
# release is a list of requests (ConnInfo objects), ordered by age, newest
# last.
gGetWorkQueue = defaultdict(deque)

# Once we've matched a 'pid' and a 'get_work' the connection from the client
# goes in here (the connection from the worker is closed once we've sent it the
# work). The key here is the unique 'tag' we gave the worker that it'll give us
# back when it makes its 'return_results' call. We don't store the release
# here, that match has already been succesfully made.
gWorkInProgress = {}

# File numbers to sockets
gClientMap = {}

# Data for get_tag()
gSeqNo = 0

gErrCount = 0

gBusyTime = 0
gIdleTime = 0

gIdList = defaultdict(lambda: 0)

# Map from connection's ID to last time received
gClientList = {}
gWorkerList = {}

# -----------------------------------------------------------------------------

# Poor man's UUID
def get_tag():
    global gSeqNo

    n = gSeqNo
    gSeqNo += 1

    return str(int(time.time())) + '.' + str(os.getpid()) + '.' + str(n)

# -----------------------------------------------------------------------------

def is_complete_query(msg):
    # headers are always terminated by \r\n\r\n
    kTerminator = '\r\n\r\n';

    if kTerminator not in msg: return False

    headerSize = msg.find(kTerminator)+4

    cl = 'Content-Length: '
    # No content-length, therefore no body
    if cl not in msg: return True
    lpos = msg.find(cl)

    # Parse the context-length header
    newline = msg.find('\r', lpos+len(cl))
    bodyLength = int(msg[lpos+len(cl):newline])

    return len(msg) == headerSize+bodyLength

# -----------------------------------------------------------------------------

def ok200():
    return 'HTTP/1.0 200 OK\r\n'\
           'Server: LEMServer/0.1.0\r\n'\
           'Content-Type: text/plain\r\n'\
           '\r\n'

# -----------------------------------------------------------------------------

def ok200_html():
    return 'HTTP/1.0 200 OK\r\n'\
           'Server: LEMServer/0.1.0\r\n'\
           'Content-Type: text/html\r\n'\
           '\r\n'

# -----------------------------------------------------------------------------

def unimp501():
    return 'HTTP/1.0 501 Not Implemented\r\n'\
           'Server: LEMServer/0.1.0\r\n'\
           'Content-Type: text/plain\r\n'\
           '\r\n'\
           "I don't know how to do that\r\n"

# -----------------------------------------------------------------------------

gLastSeqNo = 0
gLastStatus = time.time()
def status_text():
    global gLastSeqNo, gLastStatus

    rate = int((gSeqNo-gLastSeqNo)/(time.time()-gLastStatus))
    gLastSeqNo = gSeqNo
    gLastStatus = time.time()

    msg = time.strftime('%Y-%m-%d %H:%M:%S %Z: ')
    msg += 'Handled ' + str(gSeqNo) + ' queries ('+str(rate)+' Hz). '
    msg += str(len(gWorkInProgress)) + ' workers working, '
    msg += str(nConns(gGetWorkQueue)) + ' waiting, '
    msg += str(nConns(gPIDQueue)) + ' backlog. '
    msg += str(gErrCount) + ' errors. '
    msg += 'Server '+str(.1*int(1000*gBusyTime/(gBusyTime+gIdleTime)+.5))+'% busy. '

    # Forget about clients and workers after 10 minutes
    t0 = time.time()
    n10s = 0
    n60s = 0
    n5m = 0
    n30m = 0
    for k in gClientList.keys():
        t = gClientList[k]
        if t0 < t+10: n10s += 1
        if t0 < t+60: n60s += 1
        if t0 < t+300: n5m += 1
        if t0 < t+1800: n30m += 1
        if t0 > t+1800: gClientList.pop(k)
    for k in gWorkerList:
        if t0 > gWorkerList[k] + 600: gClientList.pop(k)

    msg += str(n10s)+'/'+str(n60s)+'/'+str(n5m)+'/'+str(n30m)+' clients active in last 10s/60s/5m/30m. '

    msg += str(len(gWorkerList)) + ' workers in last 10m.'

    return msg

# -----------------------------------------------------------------------------

def handle_status(client, args):
    return {client: ok200()+status_text()+'\n'}

# -----------------------------------------------------------------------------

def handle_releases(client, args):
    rels = set()
    # TODO ideally we would also list releases with only in-progress PIDs,
    # but we don't store that information anywhere.
    for key in gPIDQueue:
        if len(gPIDQueue[key]) > 0:
            rels.add(key)

    msg = ok200()
    for rel in rels: msg += rel+'\n'
    return {client: msg}

# -----------------------------------------------------------------------------

def handle_backlog(client, args):
    if not 'rel' in args:
        count = {}
        for rel in gPIDQueue.keys():
            N = len(gPIDQueue[rel])
            if N > 0: count[rel] = N

        msg = ok200()
        if len(count) == 0:
            msg += 'No backlog\n'
        else:
            msg += 'Backlogs:\n'
            for rel in count.keys():
                msg += rel+': '+str(count[rel])+'\n'
        return {client: msg}

    # Otherwise a release /was/ specified, return machine-readable result
    rel = failsafe_arg('backlog', args, 'rel')

    count = len(gPIDQueue[rel])

    return {client: ok200()+str(count)+'\n'}

# -----------------------------------------------------------------------------

def handle_ids(client, args):
    msg = ok200()

    for key in sorted(gIdList.keys()):
        msg += key + ' : ' + str(gIdList[key])+'\n'

    return {client: msg}

# -----------------------------------------------------------------------------

def handle_status_html(client, args):
    htmlHead = '<html><head><meta charset="UTF-8"><meta http-equiv="refresh" content="60"><title>LEM server status</title></head><body>\n'
    htmlFoot = '</body></html>\n'

    return {client: ok200_html() + htmlHead + status_text() + htmlFoot}

# -----------------------------------------------------------------------------
def failsafe_arg(cmd, args, name):
    if name in args: return args[name]
    print 'Error:', cmd, 'without', name, args
    return None

# -----------------------------------------------------------------------------

def handle_get_work(client, args):
    global gQueuedPIDRequest, gGetWorkQueue, gWorkInProgress, gVerbosity

    workerID = failsafe_arg('get_work', args, 'id')
    rel = failsafe_arg('get_work', args, 'rel')
    timeout = failsafe_arg('get_work', args, 'timeout')

    if not workerID or not rel or not timeout:
        return {client: unimp501()}

    # No work to give out right now
    if len(gPIDQueue[rel]) == 0 or gPIDQueue[rel][-1].late():
        # Experimenting with not letting too many get_work calls build up
        gGetWorkQueue[rel].append(ConnInfo(client, '', min(10,timeout)))
        return {}

    # There is work, and it's recent enough
    info = gPIDQueue[rel].pop()
    # Remember we're doing this so we know who the reply goes to
    tag = get_tag()
    gWorkInProgress[tag] = info


    if gVerbosity > 0:
        print 'Handing out work under tag', tag

    return {client: ok200() + 'tag='+tag+'\n' + info.query}

# -----------------------------------------------------------------------------

def handle_return_results(client, args, postdata):
    global gQueuedPIDRequest, gGetWorkQueue, gWorkInProgress, gVerbosity

    tag = failsafe_arg('return_results', args, 'tag')

    if not tag:
        return {client: unimp501()}

    if tag not in gWorkInProgress:
        print 'Warning: Reply for unknown tag', tag
        return {client: ok200()}

    workerClient = gWorkInProgress.pop(tag).client

    if gVerbosity > 0: print 'Sending response', postdata

    # No point checking the timeout, just try to send to the client in any case
    return {client: ok200(),
            workerClient: ok200()+postdata}

# -----------------------------------------------------------------------------

def handle_pid(client, args, postdata):
    global gQueuedPIDRequest, gGetWorkQueue, gWorkInProgress, gVerbosity

    rel = failsafe_arg('handle_pid', args, 'rel')
    timeout = failsafe_arg('handle_pid', args, 'timeout')

    if not rel or not timeout:
        return {client: unimp501()}

    info = ConnInfo(client, postdata, timeout)

    # No recent matching worker, stick it in the queue
    if len(gGetWorkQueue[rel]) == 0 or gGetWorkQueue[rel][-1].late():
        gPIDQueue[rel].append(info)
        return {}

    # All OK, hook things up
    workerClient = gGetWorkQueue[rel].pop().client

    # Remember that we're doing this
    tag = get_tag()
    gWorkInProgress[tag] = info


    if gVerbosity > 0: print 'Handing out work under tag', tag
    return {workerClient: ok200() + 'tag='+tag+'\n' + info.query}

# -----------------------------------------------------------------------------

# Returns a map from sockets to the strings to send to them
def handle_query(sockStr, client):
    if gVerbosity > 0: print sockStr

    query, junk, postdata = sockStr.partition('\r\n\r\n')

    isGet = query.startswith('GET')
    isPost = query.startswith('POST')

    # Second thing after GET/POST. Drop the leading slash
    params = query.split(' ')[1][1:]
    cmd, junk, argStr = params.partition('?')

    args = {}
    for arg in argStr.split('&'):
        key, junk, val = arg.partition('=')
        args[key] = val

    if gVerbosity > 0: print 'Parameters are:', params
    if gVerbosity > 0: print 'Cmd:', cmd, 'args:', args

    shortId = 'unknown'
    if 'id' in args:
        shortId = args['id']

        # Update last contacted
        if isPost and cmd == 'pid':
            gClientList[shortId] = time.time()
        if isGet and cmd == 'get_work':
            gWorkerList[shortId] = time.time()

        # Clumsily truncate the last two numbers off the end to leave the host
        if shortId.rfind('.') > 0: shortId = shortId[:shortId.rfind('.')]
        if shortId.rfind('.') > 0: shortId = shortId[:shortId.rfind('.')]

        gIdList[shortId] += 1


    get_handlers = {
        'get_work'    : handle_get_work,
        'status'      : handle_status,
        'status.html' : handle_status_html,
        'releases'    : handle_releases,
        'backlog'     : handle_backlog,
        'ids'         : handle_ids
    }

    post_handlers = {
        'pid'            : handle_pid,
        'return_results' : handle_return_results
    }

    if isGet and cmd in get_handlers:
        return get_handlers[cmd](client, args)

    if isPost and cmd in post_handlers:
        return post_handlers[cmd](client, args, postdata)

    # Didn't hit any of the cases above
    print 'Rejecting malformed query:'
    print sockStr

    return {client: unimp501()}

# -----------------------------------------------------------------------------

def startup(port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    # Reuse port immediately even if a previous instance just aborted.
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(('', port))
    # Parameter is the length of "backlog" allowed
    s.listen(128)
    return s

# -----------------------------------------------------------------------------

def close_old_connections():
    toClose = []
    nClient = 0
    nWorker = 0
    # Clear out old connections
    for rel in gPIDQueue.keys():
        toClose += pop_old_conns(gPIDQueue[rel])
    nClient = len(toClose)
    for rel in gGetWorkQueue.keys():
        toClose += pop_old_conns(gGetWorkQueue[rel])
    nWorker = len(toClose) - nClient

    if nClient > 0:
        print 'Closing', nClient, 'client connections that reached their timeout'

    if nWorker > 0 and gVerbosity > 0:
        print 'Closing', nWorker, 'worker connections that reached their timeout'

    for info in toClose:
        gClientMap.pop(info.client.fileno())
        info.client.close()

# -----------------------------------------------------------------------------

def shed_load():
    toClose = []

    kConnLimit = 3300

    # First drop older PID requests
    while nConns(gPIDQueue) + nConns(gGetWorkQueue) + len(gWorkInProgress) > kConnLimit and nConns(gPIDQueue) > 0:
        toClose += pop_fronts(gPIDQueue)

    # If things get desperate, drop get_work requests
    while nConns(gGetWorkQueue) + len(gWorkInProgress) > kConnLimit and nConns(gGetWorkQueue) > 0:
        toClose += pop_fronts(gGetWorkQueue)

    # And finally, throw away actual work in progress (at random)
    keys = gWorkInProgress.keys()
    for key in keys:
        if len(gWorkInProgress) <= kConnLimit: break
        toClose.append(gWorkInProgress.pop(key))


    if len(toClose) > 0:
        print 'Overloaded: closing', len(toClose), 'connections'
        # Doesn't seem to recover situation. Let the restarter do its thing
        sys.exit(1)

    for info in toClose:
        gClientMap.pop(info.client.fileno())
        info.client.close()


# -----------------------------------------------------------------------------

if __name__ == '__main__':
    kMinPort = 1024
    kMaxPort = 1031
    kNumPorts = kMaxPort-kMinPort+1

    serverSock = [startup(port) for port in range(kMinPort, kMaxPort+1)]

    print 'Listening on ports', kMinPort, '-', kMaxPort

    ep = select.epoll()
    socketMap = {}
    for ss in serverSock:
        ep.register(ss, select.POLLIN)
        socketMap[ss.fileno()] = ss

    inbox = defaultdict(lambda: '')
    outbox = defaultdict(lambda: '')

    statusTime = time.time()

    t1 = time.time()
    t2 = time.time()

    while True:
        # No more than once every three seconds
        if time.time() - statusTime > 3:
            statusTime = time.time()
            print status_text()
            # Busy fraction averages between these updates
            gBusyTime = 0
            gIdleTime = 0

        # Do some housekeeping
        close_old_connections()
        shed_load()

        # Update status every minute even if nothing's happening
        t1 = time.time()
        gBusyTime += t1-t2
        events = ep.poll(60)
        t2 = time.time()
        gIdleTime += t2-t1

        for event in events:
            fileno, mask = event

            if fileno not in socketMap and fileno not in gClientMap:
                print 'Unknown fileno!', fileno

            if fileno in socketMap:
                assert mask == select.POLLIN

                ss = socketMap[fileno]
                clientSock, addr = ss.accept()

                clientSock.setblocking(0)

                # We'll always start by wanting to read this socket
                ep.register(clientSock, select.POLLIN)
                gClientMap[clientSock.fileno()] = clientSock

            if fileno in gClientMap:
                cs = gClientMap[fileno]

                if mask == select.POLLIN:
                    x = cs.recv(4096)
                    if x == '':
                        print 'Read nothing. Client hung up on us?'
                        gErrCount += 1
                        gClientMap.pop(fileno)
                        cs.close()
                    else:
                        inbox[cs] += x
                        if is_complete_query(inbox[cs]):
                            # Will never need to read again
                            ep.unregister(cs)
                            msgs = handle_query(inbox[cs], cs)
                            inbox.pop(cs)
                            for sock in msgs:
                                outbox[sock] = msgs[sock]
                                ep.register(sock, select.POLLOUT)

                elif mask == select.POLLOUT:
                    msg = outbox[cs]
                    n = cs.send(msg)
                    if n == 0:
                        print 'Wrote nothing. Client hung up on us?'
                        gErrCount += 1
                        gClientMap.pop(fileno)
                        cs.close()
                    else:
                        msg = msg[n:]
                        outbox.pop(cs)
                        if msg == '':
                            # Sent everything
                            gClientMap.pop(fileno)
                            cs.close()

                else:
                    print 'Unknown mask', mask
                    gClientMap.pop(fileno)
                    cs.close()
                    gErrCount += 1