loader.py

#!/usr/bin/env python

'''greenhouse network loader

this script's only goal is to get a greenhouse-based server up and running and
IO bound.
'''

import optparse
import socket
import sys

from greenhouse import io, pool, scheduler, utils


DEFAULT_ADDRESS = "0.0.0.0"
DEFAULT_PORT = 9000
DEFAULT_CONCURRENCY = 4096


def server_connection(clientsock):
    clientfile = clientsock.makefile()

    ping = clientsock.recv(8192)
    if ping.strip().lower() != "ping":
        clientsock.close()
        return

    clientsock.sendall("PONG")

    rally = clientsock.recv(8192)
    if rally.strip().lower() != "rally's":
        clientsock.close()
        return

    clientsock.sendall("ON")
    clientsock.close()

def server(options):
    print "server starting on address %s port %d." % (
            options.address, options.port)
    print "shut it down with <Ctrl>-C"
    try:
        serversock = io.Socket()
        serversock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        serversock.bind((options.address, options.port))
        serversock.listen(socket.SOMAXCONN)
        while 1:
            clientsock, address = serversock.accept()
            scheduler.schedule(server_connection, args=(clientsock,))
    except KeyboardInterrupt:
        print "KeyboardInterrupt caught, closing listener socket"
        serversock.close()


def client_connection(options, p, e):
    try:
        sock = io.Socket()
        sock.connect((options.address, options.port))

        sock.sendall("PING")

        pong = sock.recv(8192)
        if pong.strip().lower() != "pong":
            sock.close()
            p.put(options, p, e)
            return

        sock.sendall("RALLY'S")

        on = sock.recv(8192)
        sock.close()
        p.put(options, p, e)
    except KeyboardInterrupt:
        e.set()

def client(options):
    concurrency = options.concurrency
    p = pool.OneWayPool(client_connection, options.concurrency)
    p.start()
    e = utils.Event()
    for i in xrange(options.concurrency):
        p.put(options, p, e)
    try:
        e.wait()
    except KeyboardInterrupt:
        pass


def main():
    parser = optparse.OptionParser()
    parser.add_option("-a", "--address", default=DEFAULT_ADDRESS)
    parser.add_option("-p", "--port", type="int", default=DEFAULT_PORT)
    parser.add_option("-c", "--concurrency", type="int",
            default=DEFAULT_CONCURRENCY)

    options, args = parser.parse_args()

    {'server': server, 'client': client}[args[0]](options)


if __name__ == '__main__':
    main()