#!/bin/bash


backtrace() {

        local funcname="" sourcefile="" lineno="" n e s="yes" args newarg

        declare -i strip=1

        if [[ -n $1 ]]; then
                strip=$(( $1 ))
        fi

        echo "Call stack: (most recent first)"
        for (( n = $strip, p = 0 ; n < ${#FUNCNAME[@]} ; ++n )) ; do
                funcname=${FUNCNAME[$n]}
                sourcefile=$(basename ${BASH_SOURCE[$n]})
                lineno=${BASH_LINENO[$n - 1]}
                # Display function arguments
                if [[ ! -z "${BASH_ARGV[@]}" ]]; then
                        args=
                        for (( j = ${BASH_ARGC[$n - 1]}; j > 0; j-- )); do
                                newarg=${BASH_ARGV[$j + $p - 1]}
                                args="${args:+${args} }'${newarg}'"
                        done
                        (( p += ${BASH_ARGC[$n - 1]} ))
                fi
                echo "  ${funcname} ${args:+${args} }at ${sourcefile}:${lineno}"
        done

}

trap '(echo "Untrapped error"
echo
backtrace
echo "Something bad happened, please fix & retry"
echo "Exiting!"
)' ERR

set -e

export LST_SESSION=$$

# Number of threads to use to send requests
# NB: this takes the test and * NCONN - so 1MB loop 500 concurrnecy=8 means 8 threads doing 500 1MB.
MAXNET=${MAXNET:-1}
NCONN=${NCONN:-1}
DISTR=${DISTR:-"1:1"}
DIRECTION=${DIRECTION:-"--from clients --to servers"}

for var in CLIENTS SERVERS; do
	if [ -z "${!var}" ]; then
		echo "variable $var is not set!"
		exit 1
	fi
done

export CLIENTS SERVERS

host2gni() {
	local host_name=$1
	echo "$(grep ${host_name} /etc/hosts | awk '{print $1}')@gni"
}

# setup requires the environ variables CLIENTS and SERVERS to be set
#  with donut hostnames
setup() {

	if [ -z "$CLIENTS" -o -z "$SERVERS" ]; then
		echo "Please set CLIENTS and SERVERS with space delimited hostnames"
		return 1
	fi

	mark_logs starting new LST session $LST_SESSION
	# do we want to use --force to make this more robust ?
	lst new_session --force test.$(date +%s)

	SERVER_NIDS=""
	CLIENT_NIDS=""

	for server in $SERVERS; do
		if echo $server | grep -q '@'; then
			SERVER_NIDS="$SERVER_NIDS $server"
		else
			SERVER_NIDS="$SERVER_NIDS $(host2gni $server)"
		fi
	done
	for client in $CLIENTS; do
		if echo $client | grep -q '@'; then
			CLIENT_NIDS="$CLIENT_NIDS $client"
		else
			CLIENT_NIDS="$CLIENT_NIDS $(host2gni $client)"
		fi
	done

	if [ -z "$SERVER_NIDS" -o -z "$CLIENT_NIDS" ]; then
		echo "Translating $SERVERS or $CLIENTS failed: '$SERVERS' or '$CLIENTS'"
		return 1
	fi

	mark_logs adding group servers $SERVER_NIDS
	lst add_group servers $SERVER_NIDS
	lst list_group servers --all
	mark_logs adding group clients $CLIENT_NIDS
	lst add_group clients $CLIENT_NIDS
	lst list_group clients --all
	lst add_batch ping_test
	lst add_batch write_test
	lst add_batch read_test
}

cleanup () {
	lst end_session
	return 0
}

mark_logs() {
	msg=$*

	local s=${SERVERS// /,}
	local c=${CLIENTS// /,}
	#pdsh -S -w $c,$s lctl mark $msg
	lctl mark $msg
}

setup_ping_test() {
	local size=$1
	# we ignore size, just here for a common api for write/read
	local pings=$2
	local testname=${3:-"ping_test"}

	mark_logs configuring $testname $pings:$NCONN with LST_SESSION=$LST_SESSION
	bash -x -c "lst add_test --batch $testname --distribute ${DISTR} --concurrency $NCONN $DIRECTION --loop $pings ping"
	lst query $testname --all
}

setup_write_test() {
	local size=$1
	local nloop=${2:-"-1"}
	local testname=${3:-"write_test"}

	mark_logs configuring $testname loop=$nloop size=$size with LST_SESSION=$LST_SESSION
	bash -x -c "lst add_test --batch $testname --distribute ${DISTR} --concurrency $NCONN --loop $nloop $DIRECTION brw write check=full size=$size"
	lst query $testname --all
}


setup_read_test() {
        local size=$1
	local nloop=${2:-"-1"}
	local testname=${3:-"read_test"}

        mark_logs configuring $testname loop=$nloop size=$size with LST_SESSION=$LST_SESSION
        bash -x -c "lst add_test --batch $testname --distribute ${DISTR} --concurrency $NCONN --loop $nloop $DIRECTION brw read check=full size=$size"
        lst query $testname --all
}

watch_test() {
	local testname=$1
	local duration=${2:-0}
	local now=0
	local begin=$(date +%s)
	local trigger=1

	now=$begin
	set +e
	while true; do
		lst query $testname | grep -q "running"
		if [ ${PIPESTATUS[1]} -ne 0 ]; then
			break
		fi
		# if no duration, we'll wait until it is done
		if [ $duration -ne 0 -a $((now - begin)) -ge $duration ];then
			echo "test $testname ran out of time!"
			break
		fi
		sleep 1
		let now=$((now + 1))
	done

	lst stop $testname || true
	echo "test $testname is done!"
}

run_test() {
	local testname=$1
	local duration=$2
	local interval=5
	local mpid

	#lctl clear
	mark_logs "starting $testname with LST_SESSION=$LST_SESSION"
	lst query $testname --all
	lst list_batch $testname

	# it would be better to try and use explicit nids sometimes
	# perhaps CLIENT_NIDS from setup() ?
	lst stat clients servers --delay $interval &
	mpid=$!

	lst run $testname
	watch_test $testname $duration

	kill -term $mpid
	lst query $testname --all
	lst list_batch $testname
	lst show_error clients servers
	mark_logs "ending $testname with LST_SESSION=$LST_SESSION"
}

do_it() {
	local name=$1
	local size=$2
	local nloops=$3

	setup_${name}_test $size $nloops
	time1=$(date +%s.%N)
        run_test ${name}_test
	time2=$(date +%s.%N)
	diff=$(echo "$time2 - $time1" | bc -l)
	printf "test '$name $size $nloops' on ${NCONN} thr session=$LST_SESSION done in %.04fs\n" ${diff}

	case $name in
		ping)  
			total_pings=$((MAXNET *NCONN * $nloops))
			rate=$(echo "$total_pings / $diff" | bc -l)
			printf "completed $total_pings pings in %.03fs\n" ${diff}
			printf "rate: %.03f pings/s\n" ${rate}
			;;
		write|read)
			parse_size=$(echo $size | grep -o "[0-9]*")
			ks=$(echo $size | grep -o [kK])	
			ms=$(echo $size | grep -o [mM])	
			extra=1
			txt="B"
			if [ -n "$ks" ]; then
				extra=1024
				txt="KB"
			elif [ -n "$ms" ]; then
				extra=$((1024*1024))
				txt="MB"
			fi
			parse_size=$((parse_size * extra))
			total_size=$((MAXNET * NCONN * $nloops * parse_size))
			rate_size=$(($total_size / $extra))
			rate=$(echo "($total_size / $diff) / $extra" | bc -l)
			printf "completed $name of ${rate_size}${txt} in %.03fs\n" ${diff}
			printf "rate: %.03f ${txt}/s\n" ${rate}
			;;
	esac
}