#!/usr/bin/python
import threading
import random
import sys
import time
import os
import signal
import select
import cmd

options = {
	"filelist": [],
	"threads": 10,  # initial number of threads
	"blocksize": 40960, # amount of data (in bytes) read at a time
	"bitrate": 2*1024*1024, # target bitrate (in bits/second) for each thread
	"statinterval": 1, # time in seconds between two statistic displays
	"directio": 0, # use O_DIRECT to bypass buffercache (avoid "cheating")
	"sleepdelay": 0.1, # a thread which doesn't need data will sleep that much
	                   # seconds before checking again
	"benchduration": 0, # in seconds ; 0 = infinite
	"quiet": 0, # to suppress statistics display (enable again by pressing ENTER)
	"spawnthreshold": 100*1024, # if buffer size is less than that, create a new thread
	"killthreshold": 4*1024*1024, # if buffer size is more than that, kill a thread
	}

### WARNING : with directio, blocksize should be a multiple of the underlying device blocksize
### (with kernel 2.6, a multiple of 512 suffice, see open(2))
### WARNING : with directio, buffers should be aligned, but python doesn't ensure this, so
### directio is actually impossible until someone writes a patch or an extension for python.

class reader (threading.Thread):
	def __init__ (self):
		threading.Thread.__init__ (self)
		self.bytestransferred = 0
		self.lastbytestransferred = 0
		self.starttime = time.time () + random.random () * self.period ()
		self.running = 1
		self.nextdeadline = self.starttime
		self.filedesc = -1
		self.nextfile ()
		self.start ()
	def nextfile (self):
		if self.filedesc >= 0 : os.close (self.filedesc)
		self.filename = random.choice (options["filelist"])
		flags = os.O_RDONLY
		if options["directio"]: flags |= os.O_DIRECT
		self.filedesc = os.open (self.filename, flags)
		self.filesize = os.lseek (self.filedesc, 0, 2)
		self.filepos  = int (random.random()*self.filesize)
		self.filepos -= self.filepos % options["blocksize"] #ensure alignment
		os.lseek (self.filedesc, self.filepos, 0)
		#print ("Selected file %s (FD=%d), at offset %d/%d"
		#      %(self.filename,self.filedesc,self.filepos,self.filesize))
	def period (self):
		return 8.0*options["blocksize"]/options["bitrate"]
	def run (self):
		while self.running:
			while time.time () < self.nextdeadline:
				time.sleep (options["sleepdelay"])
			self.nextdeadline += self.period ()
			l = len (os.read(self.filedesc,options["blocksize"]))
			self.filepos += l
			if l==0: self.nextfile ()
			self.bytestransferred += l

class console (threading.Thread, cmd.Cmd):
	def __init__(self):
		threading.Thread.__init__(self)
		cmd.Cmd.__init__(self)
		self.eofcount=2
		self.updateprompt ()
		self.start ()
	def updateprompt (self):
		if options["quiet"]: self.prompt = "StreamBench> "
		else: self.prompt = ""
	def run (self):
		while running:
			try: self.cmdloop ()
			except Exception,e:
				print e
	def precmd (self,x):
		if not running: raise SystemExit
		return x
	def emptyline (self):
		if options["quiet"]: options["quiet"] = 0
		else: options["quiet"] = 1
		self.updateprompt ()
	def do_quit (self,x=""):
		quit ()
	def do_more (self,x):
		try: n=int(x)
		except: n=1
		print "Starting %d extra thread(s)."%n
		options["threads"]+=n
	def do_less (self,x):
		try: n=int(x)
		except: n=1
		print "Stopping %d thread(s)."%n
		options["threads"]-=n
	def do_EOF (self,x):
		if self.eofcount:
			print "EOF Detected, reopening (hit Ctrl+D again to quit)"
			sys.stdin = open("/dev/tty")
			self.eofcount -=1
		else:
			print "EOF Detected, quitting"
			self.do_quit ()

threadpool = []
totalbytestransferred = 0
lasttotalbytestransferred = 0
starttime = time.time ()
overallworstbuffersize = 0
lasttime = 0
running = 1

###
### OPTIONAL SETTINGS
###

if sys.argv[1:]:
	exec(urllib.urlopen(sys.argv[1]).read(),{},options)

if not options["filelist"]:
	options["filelist"] = [i.strip() for i in sys.stdin.readlines () if i.strip()]

def quit (signum=-1,frame=None):
	global running
	running = 0
		
myconsole=console()

while len(threadpool) or running:
	try:
		if not running: options["threads"] = 0
		delta = options["threads"] - len(threadpool)
		if delta>0:
			print "Increasing thread pool to %d."%options["threads"]
			for i in range(delta): threadpool.append (reader())
		if delta<0:
			print "Decreasing thread pool to %d."%options["threads"]
			for i in range(abs(delta)): threadpool[i].running = 0
			for i in range(abs(delta)):
				threadpool[0].join ()
				del threadpool[0]

		now = time.time ()
		currentworstbuffersize = 0
		for t in threadpool:
			totalbytestransferred += t.bytestransferred - t.lastbytestransferred
			t.lastbytestransferred = t.bytestransferred
			buffersize = (now - t.nextdeadline) * options["bitrate"] / 8
			if buffersize > currentworstbuffersize:
				currentworstbuffersize = buffersize
			if currentworstbuffersize > overallworstbuffersize:
				overallworstbuffersize = currentworstbuffersize

		if lasttime and not options["quiet"]:
			print ("Threads : %4d @ %6d Kb/s each, total : %8d KB/s"
			       %(len(threadpool), options["bitrate"]/1024, len(threadpool)*options["bitrate"]/8192)
			       +" | Current speed : %8d KB/s | Overall speed : %8d KB/s"
			       %((totalbytestransferred-lasttotalbytestransferred)/(now-lasttime)/1024,
				 totalbytestransferred/(now-starttime)/1024)
			       +" | Largest buffer this period : %5d KB | Overall : %5d KB"
			       %(currentworstbuffersize/1024,overallworstbuffersize/1024))
		lasttime = now
		lasttotalbytestransferred = totalbytestransferred

		if currentworstbuffersize < options["spawnthreshold"]:
			options["threads"]+=1
		if currentworstbuffersize > options["killthreshold"]:
			options["threads"]-=1

		time.sleep (options["statinterval"])

		if options["benchduration"]:
			if options["benchduration"] + starttime > now:
				running = 0
				
	except Exception,e:
		print e
		running = 0

print "Waiting for console thread (press enter if it seems stuck)"
myconsole.join()
print "Bye."