#!/usr/bin/env python

# Copyright C 2002-2005 by Isilon Systems, Inc. All rights not expressly
# granted herein are reserved.
# Created in the United States of America.
# 
# This digital media is protected by U.S. and international copyright and
# intellectual property laws. Unless otherwise specified, all information
# and screens appearing as part of this digital medium, including
# software, services, documents, text, images, icons, and logos design;
# the selection, assembly, arrangement, and design thereof; and the code
# that enables its presentation, are the sole property of Isilon Systems,
# Inc.
# 
# Except as otherwise required by applicable law, no part of any
# copyrighted material may be reproduced, modified, translated,
# distributed, published, or transmitted by any means, electronic or
# mechanical, for any purpose, without the prior written consent of Isilon
# Systems. 
#  
# Isilon and Isilon Systems are registered trademarks of Isilon Systems,
# Inc.  Legion is a trademark of Isilon Systems, Inc.  All other
# trademarks and registered trademarks are the property of their
# respective holders.

import sys, os, popen2, re, time, getopt
import isi.app.lib.cluster as cluster
import isi.net

#--------------------
# Config constants
#
ttcp_path = 		"/usr/bin/isi_hwtools/isi_ttcp"
for_array_path = 	"/usr/bin/isi_for_array"
test_thread_path = 	"/usr/bin/isi_test_thread"

# (lazy_send_cq enabled, send_poll_threshold, snd size)
ttcp_runs = [(0, 0, 		2**13),
	(1, 0, 		2**13),
	(0, 0, 		2**8),
	(1, 0, 		2**8)]

latency_runs = [(0, 0),
		(1, 0)]

# Each final data point will be an average of this many runs
ttcp_num_measurements = 	10 

# Total num bytes to send for each ttcp run  (1/2 GB)
ttcp_num_bytes = 		2**29

# How many round trips for each measurement
# too big a number can cause splits
latency_test_trips =		50	

# Average this many latency tests together
latency_test_measurements =	20 

#
# End config section
#--------------------



def vprint(s) :
	if debug_output:
		print s

class TtcpResult:
	def __init__(s) :
		s.total_bytes = 	0L
		s.mbit_cpu_sec =	0.0
		s.mbit_real_sec =	0.0

		s.bytes_re = re.compile(".*ttcp-t: (\d+)\sbytes processed.*")
		# ttcp-t:  0.007941 CPU sec  = 16118.88 Mbit/cpu sec
		s.cpu_re = re.compile(
		"^.*ttcp-t\:\s+([0-9\.]+)\s+CPU sec\s*=\s*([0-9\.]+)\s+Mbit.*")
		s.real_re = re.compile(
		"^.*ttcp-t\:\s+([0-9\.]+)\s+real sec\s*=\s*([0-9\.]+)\s+Mbit.*")

	def __str__(s) :
		out = ""
		out += "total_bytes %d " % s.total_bytes
		out += "mbit_per_real_sec %f " % s.mbit_real_sec
		out += "mbit_per_cpu_sec %f" % s.mbit_cpu_sec
		return out

	def from_stream(s, sout) :

		for l in sout.readlines() :
			m = s.bytes_re.match(l)
			if m :
				s.total_bytes = long(m.group(1))
				continue

			m = s.real_re.match(l)
			if m :
				#s.real_sec = float(m.group(1))
				s.mbit_real_sec = float(m.group(2))
				continue

			m = s.cpu_re.match(l)
			if m :
				#s.cpu_sec = float(m.group(1))
				s.mbit_cpu_sec = float(m.group(2))
				continue

	def to_list(s) :
		return [s.total_bytes, s.mbit_cpu_sec, s.mbit_real_sec]

	def from_list(s, l) :
		[s.total_bytes, s.mbit_cpu_sec, s.mbit_real_sec] = l


class TtcpStat :
	def __init__(s, name = "") :
		s.name = name	
		s.min = TtcpResult() 
		s.min.from_list([sys.maxint, 999999999.0, 999999999.0])
		s.max = TtcpResult() 
		s.total = TtcpResult()
		s.N = 0
	
	def __str__(s) :
		o = ""
		if s.name :
			o += s.name + ": "
		o += "min: " + str(s.min) 
		o += "\nmax: " + str(s.max) 
		o += "\nave: " + str(s.get_average())
		return o

	def get_average(s) :
		if s.N :
			ave = TtcpResult()
			tl = s.total.to_list()
			vprint("total list: " + str(tl))
			al = map(lambda x: x/s.N, tl)
			vprint("ave list: " + str(al))

			ave.from_list(al)
#			ave.from_list(map(lambda x: x/s.N, s.total.to_list())) 
			return ave

	def add(s, result) :
		if not result:
			return
		l = map(lambda x,y: min (x,y), s.min.to_list(), result.to_list())
		s.min.from_list(l)  
		
		l = map(lambda x,y: max (x,y), s.min.to_list(), result.to_list())
		s.max.from_list(l)

		l = map(lambda x,y: x + y, s.total.to_list(), result.to_list())
		s.total.from_list(l)
		s.N += 1


class IsiIbTest :

	def __init__(s) :
		s.remote_devid = -1
		s.cluster = cluster.Cluster()
		s.local_devid = s.cluster.local()
		s.devices = s.cluster.devices()
		s.child_processes = []
		s.remote_addr_str = ""

	def usage(s) :
		print "Usage:  isi_ib_test [-n <dest-devid>] [-v]"
		print "\nisi_ib_test runs two types of performance tests:"
		print "\t- Kernel-kernel unidirectional throughput" 
		print "\t- RBM round-trip latency"
		print "Multiple runs will be performed for each data point, and "
		print "the min, max, and average values reported.  Parameters "
		print "can be changed by editing the beginning of this script.\n"
		print "The local node will be used as the sender.  The receiver"
		print "node can be specified via -n, or the first available "
		print "other node will be used.\n"
		print "The -v option turns on verbose debugging output."
		print "\nDepends on: isi_ttcp, bw_test_sdp.ko, isi_for_array"
		sys.exit(1)
		
	# Parse command args
	def process_args(s) :
		global debug_output

	    	try:
			opts, args = getopt.getopt(sys.argv[1:], "vn:")
		except getopt.GetoptError:
	    		s.usage()
	
		debug_output = 0
	   	for (o,a) in opts :
			if o == "-v" :
				debug_output = 1
			elif o == "-n" :
				s.remote_devid = int(a)
			else :
				s.usage()

		# pick first available remote device
		if (s.remote_devid == -1) :
			for d in s.devices :
				if d != s.local_devid :
					s.remote_devid = int(d)
					vprint ("# Using %d as remote "
						"device id" % d)
					break
			if s.remote_devid == -1 :
				print "! not enough devices !"
				s.usage()

		elif (s.remote_devid == s.cluster.local()) :
			print ("! devid %d is the local device !" %
				s.remote_devid)
			s.usage()

		# now we have remote devid, get ip address
		s.remote_addr_str = cluster.Device(d).ipaddress(
					isi.net.NET_INTERNAL)  
		
		if s.remote_addr_str == "" :
			print "! could not get remote ip address !"
			s.usage()

#	def device_command(s, devid, command, block = 1) :
#		if (devid == s.local_devid) :
#			if (block) :
#				os.system(command)
#			else :
#				os.popen2	

	def load_modules(s) :
		s.quiet_command("isi_for_array -Q -n%d,%d kldload bw_test_sdp"
			% (s.local_devid, s.remote_devid))	
			
		s.quiet_command("isi_for_array -Q -n%d,%d killall isi_ttcp"
			% (s.local_devid, s.remote_devid))	

	def unload_modules(s) :
		s.quiet_command("isi_for_array -Q -n%d,%d kldunload bw_test_sdp"
			% (s.local_devid, s.remote_devid))	

	def quiet_command(s, cmd) :
		obj = popen2.Popen4(cmd)
		status = obj.wait()
		if os.WIFEXITED(status) :
			error = os.WEXITSTATUS(status)
			if (error and debug_output) :
				print ("! command \"%s\" exit status %d !" %
					(cmd, error))
				vprint("### output:")
				for l in obj.fromchild.readlines() :
					vprint("### " + l)

	def ttcp_run(s, lazy_cq, poll_thresh, s_size, last_settings = False) :

		num_bufs = ttcp_num_bytes / s_size 

		if (not last_settings) :
			# set lazy_cq params, really only needed on sender
			vprint ("setting lazy_send_cq %d and poll_thresh %d"
				 % (lazy_cq, poll_thresh))

			s.quiet_command("sysctl net.ib.sdp.lazy_send_cq=%d"
				% lazy_cq)

			s.quiet_command("sysctl net.ib.sdp.send_poll_thresh=%d"
				% poll_thresh)

		# open remote receiver 
		cmd = [for_array_path, "-n" + str(s.remote_devid),
			ttcp_path, "-r -i -s -k -l8192"]  
		ctxt = " ".join(cmd)
		reader = popen2.Popen4(ctxt)
		vprint ("# CMD (%d): \"%s\"" % (reader.pid, ctxt))
	
		# wait for receiver to set up 
		time.sleep(2.5)

		status = reader.poll()
		if (status != -1) :
			print "Remote exit status %d" % status
			print reader.fromchild.readlines()
			sys.exit(1)
			
		# run local sender
		cmd = [ttcp_path, "-t -i -s -k", "-l%d" % s_size,
			"-n%d" % num_bufs, s.remote_addr_str]  
		ctxt = " ".join(cmd)
		writer = popen2.Popen4(ctxt)
		vprint("# CMD (%d): \"%s\"" % (writer.pid, ctxt))
		
		stats = TtcpResult()
		stats.from_stream(writer.fromchild)
		vprint(stats)

		vprint("# Waiting for writer")
		error = 0
		status = writer.wait()
		if os.WIFEXITED(status) :
			error = os.WEXITSTATUS(status)
			if (error) :
				print ("! ttcp exit status %d !" %  error)
				vprint("### reader:")
				for l in reader.fromchild.readlines() :
					vprint("### " + l)
		elif os.WIFSIGNALED(status) :
			print ("# pid %d signalled" % writer.pid)
			error = 5
		else :
			print "! other wait condition !"
			error = 6

		if (error) :
			return None
		else : 
			return stats

	def run_ttcp_tests(s) :

		data = map(lambda x: None, ttcp_runs)
		for i in range(0, len(ttcp_runs)) :
			(lazy_cq, poll_thresh, s_size) = ttcp_runs[i] 

			data[i] = TtcpStat() 

			for n in range(0, ttcp_num_measurements) :
				stats = s.ttcp_run(lazy_cq, poll_thresh, s_size,
					(n != 0))
				if stats :
					data[i].add(stats)
				else :
					print("Error, aborting ttcp tests")
					return 1

		# report results
		for i in range(0, len(ttcp_runs)):
			print(("### TTCP: lazy_send_cq %d poll_thresh %d " +
			"snd_size %d") % ttcp_runs[i])
			print data[i] 

	def add_latency_stat(s, stat, val) :
		(smin, smax, stot, snum) = stat
		stat = (min(smin, val), max(smax, val), stot + val, snum + 1)
		return stat

	def print_latency_stat(s, stat) :
		(smin, smax, stot, snum) = stat
		if (snum == 0) :
			print "! no measurements !"
			ave = 0
		else :
			ave = stot / float(snum)
		print("min %d max %d ave %f N %d" %  (smin, smax, ave, snum))

	def print_cq_settings(s, lazy_send_cq, poll_thresh) :
		out = ""
		if (lazy_send_cq) :
			out += ("lazy_send_cq, poll_thresh %d " % poll_thresh)
		print out

	def run_latency_tests(s, lazy_send_cq = 0, poll_thresh = 0) :

		online_devs = s.cluster.online(cluster.ONLINE_IFS)
		vprint("online devs: " + str(online_devs))
		if ((not s.local_devid in online_devs) or 
		    (not s.remote_devid in online_devs)) :
		    	print("! Devices %d and %d must mount /ifs !" 
				% (s.local_devid, s.remote_devid))
			return 1

		vprint ("setting lazy_send_cq %d and poll_thresh %d"
			% (lazy_send_cq, poll_thresh))

		s.quiet_command("sysctl net.ib.sdp.lazy_send_cq=%d"
			% lazy_send_cq)

		s.quiet_command("sysctl net.ib.sdp.send_poll_thresh=%d"
			% poll_thresh)
			
		cmd = test_thread_path + " \"rbm_latency(%d,%d)\""
		cmd = cmd % (s.remote_devid, latency_test_measurements)

		re_thr = re.compile(
		    "^Threaded.*, (\d+).*\((\d+) usec per RTT.*")
		re_nothr = re.compile(
		    "^Non-threaded.*, (\d+).*\((\d+) usec per RTT.*")

		# Each stat is min, max, total, num_measurements
		nothr_1b_rtt = (sys.maxint, 0, 0, 0)
		nothr_8kb_rtt = (sys.maxint, 0, 0, 0)
		thr_1b_rtt = (sys.maxint, 0, 0, 0)
		thr_8kb_rtt = (sys.maxint, 0, 0, 0)

		for i in range(0, latency_test_measurements) :
			vprint("# " + cmd)
			# HEAD currently splits and deadlocks if
			# you don't throttle these tests some
			time.sleep(1.0)
			test = popen2.Popen4(cmd)
			for l in test.fromchild.readlines() :
				m = re_thr.match(l)
				if m :
					if m.group(1) == "1" :
						thr_1b_rtt = \
						s.add_latency_stat(thr_1b_rtt, 
							long(m.group(2)))
					elif m.group(1) == "8" :
						thr_8kb_rtt = \
						s.add_latency_stat(thr_8kb_rtt, 
							long(m.group(2)))
					else :
						print("! invalid match !")
					continue


				m = re_nothr.match(l)
				if m :
					if m.group(1) == "1" :
						nothr_1b_rtt = \
						s.add_latency_stat(nothr_1b_rtt, 
							long(m.group(2)))
					elif m.group(1) == "8" :
						nothr_8kb_rtt = \
						s.add_latency_stat(nothr_8kb_rtt, 
							long(m.group(2)))
					else :
						print("! invalid Match !")
					continue

						
		# report latency stats
		print "### LATENCY: nothread, 1B msg RTT ",
		s.print_cq_settings(lazy_send_cq, poll_thresh)
		s.print_latency_stat(nothr_1b_rtt)	

		print "### LATENCY: nothread, 8KB msg RTT ",
		s.print_cq_settings(lazy_send_cq, poll_thresh)
		s.print_latency_stat(nothr_8kb_rtt)	

		print "### LATENCY: thread, 1B msg RTT ",
		s.print_cq_settings(lazy_send_cq, poll_thresh)
		s.print_latency_stat(thr_1b_rtt)	

		print "### LATENCY: thread, 8KB msg RTT ",
		s.print_cq_settings(lazy_send_cq, poll_thresh)
		s.print_latency_stat(thr_8kb_rtt)	
			
	def main(s) :

		s.process_args()
		s.load_modules()
		s.run_ttcp_tests()

		# (lazy_send_cq, poll_thresh)
		for (lazy_cq, poll_thresh) in latency_runs:
			s.run_latency_tests(lazy_cq, poll_thresh)
		#s.unload_modules()
		sys.exit(0)


if __name__ == "__main__" :
	IsiIbTest().main()	
