#! /usr/bin/env python

# Test network throughput.
#
# Usage:
# 1) on host_A: throughput -s [port]			# start a server
# 2) on host_B: throughput -c  count host_A [port]	# start a client
#
# The server will service multiple clients until it is killed.
#
# The client performs one transfer of count*BUFSIZE bytes and
# measures the time it takes (roundtrip!).


import sys, time
from socket import *

MY_PORT = 50000 + 42

BUFSIZE = 1024


def main():
	if len(sys.argv) < 2:
		usage()
	if sys.argv[1] == '-s':
		server()
	elif sys.argv[1] == '-c':
		client()
	else:
		usage()


def usage():
	sys.stdout = sys.stderr
	print 'Usage:    (on host_A) throughput -s [port]'
	print 'and then: (on host_B) throughput -c count host_A [port]'
	sys.exit(2)


def server():
	if len(sys.argv) > 2:
		port = eval(sys.argv[2])
	else:
		port = MY_PORT
	s = socket(AF_INET, SOCK_STREAM)
	s.bind(('', port))
	s.listen(1)
	print 'Server ready...'
	while 1:
		conn, (host, remoteport) = s.accept()
		while 1:
			data = conn.recv(BUFSIZE)
			if not data:
				break
			del data
		conn.send('OK\n')
		conn.close()
		print 'Done with', host, 'port', remoteport


def client():
	if len(sys.argv) < 4:
		usage()
	count = int(eval(sys.argv[2]))
	host = sys.argv[3]
	if len(sys.argv) > 4:
		port = eval(sys.argv[4])
	else:
		port = MY_PORT
	testdata = 'x' * (BUFSIZE-1) + '\n'
	t1 = time.time()
	s = socket(AF_INET, SOCK_STREAM)
	t2 = time.time()
	s.connect((host, port))
	t3 = time.time()
	i = 0
	while i < count:
		i = i+1
		s.send(testdata)
	s.shutdown(1) # Send EOF
	t4 = time.time()
	data = s.recv(BUFSIZE)
	t5 = time.time()
	print data
	print 'Raw timers:', t1, t2, t3, t4, t5
	print 'Intervals:', t2-t1, t3-t2, t4-t3, t5-t4
	print 'Total:', t5-t1
	print 'Throughput:', round((BUFSIZE*count*0.001) / (t5-t1), 3),
	print 'K/sec.'


main()