#!/usr/bin/python

# import necessary library.
import threading
import sys
import rados
import socket
import time
import os
import getopt


# create an interface class.
class multithread (threading.Thread):
	# define init process.
	def __init__(self, threadID, name, size, number):
		threading.Thread.__init__(self)
		# the thread id.
		self.threadID = threadID
		# the thread name.
		self.name = name
		# the size of an object.
		self.size = size
		# and the thread number to pass to each thread.
		self.number = number
	def run(self):
		#print "Starting " + self.name
		# call another function to process the test, pass the thread name to it.
		performancetest(self.name, self.size, self.number)
		#print "Exiting " + self.name

# create a function to do the performance test.
def performancetest(threadName, objectsize, threadnumber):

	# add host name to each thread.
	hostname = socket.gethostname()
	objectname = hostname + threadName
	# create an empty data.
	data = ''
	# create two local file, one for write test and the other for read test.
	localreadfilename = "/tmp/localreadfileobj" + threadName
	localwritefilename = "/tmp/localwritefileobj" + threadName

	# delete local read files if exists.
	try:
		if os.path.isfile(localreadfilename):
			os.remove(localreadfilename)
	except TypeError as e:
		print hostname, ' - ', threadName, ' - ', e

	# delete local write files if exists.
	try:
		if os.path.isfile(localwritefilename):
			os.remove(localwritefilename)
	except TypeError as e:
		print hostname, ' - ', threadName, ' - ', e

	# open the local file to fill some random data.
	#openfile = open(localwritefilename,"w")
	#openfile.write("1001101011")
	#openfile.close()

	# using librados to create a cluster handle.
	try:
		cluster = rados.Rados(conffile='')
	# exit program if any error.
	except TypeError as e:
		#print hostname, ' - ', threadName, ' - Argument validation error: ', e
		raise e

	#print hostname, ' - ', threadName, " - Created cluster handle."

	# using librados to connect to ceph.
	try:
		cluster.connect()
	# exit program if any error.
	except Exception as e:
		#print hostname, ' - ', threadName, " - connection error: ", e
		raise e
		# if no errors, print the info that cluster has been connected.
	#finally:
		#print hostname, ' - ', threadName, " - Connected to the cluster."

	#print hostname, ' - ', threadName, " - Creating a context for the pool 'rbd'."
	# search the pool by given name.
	# exit if no such pool.
	if not cluster.pool_exists('rbd'):
		raise RuntimeError('No rbd pool exists')
	# if the pool exsits, open it for later use.
	ioctx = cluster.open_ioctx('rbd')

	###################################
	# remove an object from ceph.
	# 1. remove an object from ceph cluster.
	###################################
	#print hostname, ' - ', threadName, " - Removing objectfrom pool 'rbd'."
	# mark begin time
	dbegintime = time.time()
	# remove the object by object name.
	try:
		ioctx.remove_object(objectname)
	except:
		print hostname, ' - ', threadName, ' - No such object, continue'
	# mark end time
	dendtime = time.time()
	ioctx.close()

	###################################
	# write an object into ceph.
	# 1. read some data from a local file.
	# 2. save the string.
	# 3. write the data to ceph cluster.
	###################################
	#print hostname, ' - ', threadName, " - Writing object with contents to pool 'rbd'."
	# open the file to read
	#openfile = open(localwritefilename,"r")
	# read the file into variable data.
	#data = openfile.read();
	# close the file.
	#openfile.close()
	# mark begin time.
	wbegintime = time.time()
	# search the pool by given name.
	# exit if no such pool.
	if not cluster.pool_exists('rbd'):
		raise RuntimeError('No rbd pool exists')
	# if the pool exsits, open it for later use.
	ioctx = cluster.open_ioctx('rbd')
	# write data into object storage by object name.
	for i in range(0,objectsize / 10):
		data = data + '1001011010'
	ioctx.write(objectname, data)
	# mark end time.
	wendtime = time.time()
	# save the size the ceph wrote.
	#wsize = os.path.getsize(localwritefilename)
	ioctx.close()

	# clean the cache
	with open('/proc/sys/vm/drop_caches', 'w') as stream:
		stream.write('1\n')

	###################################
	# read an object from ceph.
	# 1. read some data from ceph cluster.
	# 2. save the string.
	# 3. write the data to a local file.
	###################################
	#print hostname, ' - ', threadName, " - Reading object with contents from pool 'rbd'."
	# mark begin time
	rbegintime = time.time()
	# search the pool by given name.
	# exit if no such pool.
	if not cluster.pool_exists('rbd'):
		raise RuntimeError('No rbd pool exists')
	# if the pool exsits, open it for later use.
	ioctx = cluster.open_ioctx('rbd')
	# read data from ceph by object name and size.
	# then save the string into variable data.
	data = ioctx.read(objectname,objectsize)
	# mark end time
	rendtime = time.time()
	ioctx.close()
	# open the local file to write.
	openfile = open(localreadfilename,"a")
	# write data into local file.
	openfile.write(data)
	# close the open file.
	openfile.close()
	# save the size the ceph read.
	rsize = os.path.getsize(localreadfilename)

	# close the connection to ceph cluster.
	#print hostname, ' - ', threadName, " - Closing the connection."
	#ioctx.close()

	# shutdown the cluster handle.
	#print hostname, ' - ', threadName, " - Shutting down the handle."
	cluster.shutdown()

	# print the object size that read from and wrote to ceph.
	#print hostname, ' - ', threadName, " - Size(Write):  %f ." %wsize
	#print hostname, ' - ', threadName, " - Size(Read):  %f ." %rsize
	# print the period for each operations.
	#print hostname, ' - ', threadName, " - Seconds(Write):  %f ." %(wendtime - wbegintime)
	#print hostname, ' - ', threadName, " - Seconds(Read):   %f ." %(rendtime - rbegintime)
	#print hostname, ' - ', threadName, " - Seconds(Remove): %f ." %(dendtime - dbegintime)
	#print hostname, ' - ', threadName, " - Total Seconds:   %f ." %((wendtime - wbegintime) + (rendtime - rbegintime) + (dendtime - dbegintime))
	# clean the cache
	with open('/proc/sys/vm/drop_caches', 'w') as stream:
		stream.write('1\n')

# main function
def main(argv):
	# get thread number and object size form parameters
	threadnumber = 0
	objectsize = 0

	# get thread number and object size form parameters
	try:
		opts, args = getopt.getopt(argv,"ht:s:",["threadnumber=","objectsize="])
	except getopt.GetoptError:
		# if the parameter format is wrong, print the help
		print 'c-obj.py -t <threadnumber> -s <objectsize>'
		sys.exit(2)
	for opt, arg in opts:
		if opt == '-h':
			print 'c-obj.py -t <threadnumber> -s <objectsize(byte)>'
			print 'Please limit the thread number under 1000 and object size under 1MB'
			print 'since the test server is t2.small (low performance)'
			sys.exit()
		elif opt in ("-t", "--threadnumber"):
			if int(arg) > 1000:
				print 'Please limit the thread number under 1000 and object size under 1MB'
				print 'since the test server is t2.small (low performance)'
				sys.exit(2)
			threadnumber = int(arg)
		elif opt in ("-s", "--objectsize"):
			if int(arg) > 1000000:
				print 'Please limit the thread number under 1000 and object size under 1MB'
				print 'since the test server is t2.small (low performance)'
				sys.exit(2)
			objectsize = int(arg)

	# multi threading
	for i in range(0,threadnumber):
		thr = multithread(i, "Thread-%d" %i, objectsize, threadnumber)
		thr.start()

# program entry
if __name__ == "__main__":
	main(sys.argv[1:])
