#!/bin/bash

######################################################################
# customize per survey

# CHOOSE EITHER scsidevs or rawdevs
# the SCSI devices to measure - WARNING: will be erased.
# The raw devices to use
# rawdevs=${rawdevs:-"/dev/raw/raw1"}
# scsidevs=`ls /dev/sd[a-z] /dev/sd[a-z][a-z]` # all devices, if you use udev

# result file prefix.  date/time+hostname makes unique
# NB ensure the path exists if it includes subdirs
rslt_loc=${rslt_loc:-"/tmp"}
rslt=${rslt:-"$rslt_loc/sgpdd_survey_`date +%F@%R`_`uname -n`"}

# what to do (read or write)
actions=${actions:-"write read"}

# total size per device (MBytes)
# NB bigger than device cache is good
size=${size:-8192}

# record size (KBytes)
rszlo=${rszlo:-1024}
rszhi=${rszhi:-1024}

# Concurrent regions per device
crglo=${crglo:-1}
crghi=${crghi:-256}

# threads to share between concurrent regions per device
# multiple threads per region simulates a deeper request queue
# NB survey skips over #thr < #regions and #thr/#regions > SG_MAX_QUEUE
thrlo=${thrlo:-1}
thrhi=${thrhi:-4096}

#####################################################################
# leave the rest of this alone unless you know what you're doing...

# and max # threads one instance will spawn
SG_MAX_QUEUE=16

# is the sg module loaded?
sg_is_loaded=$(grep -q "^sg " /proc/modules && echo true || echo false)

# did we load it?
sg_was_loaded=false

# map given device names into SG device names
i=0
devs=()
if [ "$scsidevs" ]; then
        # we will test for a LUN, the test for a partition
        # if the partition number is > 9 this will fail

    # make sure sg kernel module is loaded
    if ! $sg_is_loaded; then
	echo "loading the sg kernel module"
	modprobe sg && sg_was_loaded=true
	sg_is_loaded=true
    fi

    for d in $scsidevs; do
	devs[$i]=`sg_map | awk "{if (\\\$2 == \"$d\") print \\\$1}"`
        if [ -z "${devs[i]}" ]; then
            echo "Can't find SG device for $d, testing for partition"
            pt=`echo $d | sed 's/[0-9]$//'`
            # Try again
            devs[$i]=`sg_map | awk "{if (\\\$2 == \"$pt\") print \\\$1}"`
            if [ -z "${devs[i]}" ]; then
                echo -e "Can't find SG device $pt.\nDo you have the sg module configured for your kernel?"
                exit 1
           fi
	fi
	i=$((i+1))
    done
elif [ "$rawdevs" ]; then
    for r in $rawdevs; do
	RES=`raw -q $r`
	if [ $? -eq 0 ];then
	    devs[$i]=$r
	    i=$((i+1))
	else
	    echo "Raw device $r not set up"
	    exit 1
	fi
    done
else
    echo "Must specify scsidevs or rawdevs"
    exit 1
fi

ndevs=${#devs[@]}

# determine block size. This should also work for raw devices
# If it fails, set to 512
bs=$((`sg_readcap -b ${devs[0]} | awk '{print $2}'`))
if [ $bs == 0  ];then
	echo "sg_readcap failed, setting block size to 512"
	bs=512
fi
rsltf=${rslt}.summary
workf=${rslt}.detail
echo -n > $rsltf
echo -n > $workf

print_summary () {
    if [ "$1" = "-n" ]; then
	minusn=$1; shift
    else
	minusn=""
    fi
    echo $minusn "$*" >> $rsltf
    echo $minusn "$*"
}

print_summary "$(date) sgpdd-survey on $rawdevs$scsidevs from $(hostname)"

for ((rsz=$rszlo;rsz<=$rszhi;rsz*=2)); do
    for ((crg=$crglo;crg<=$crghi;crg*=2)); do 
	for ((thr=$thrlo;thr<=$thrhi;thr*=2)); do
	    if ((thr < crg || thr/crg > SG_MAX_QUEUE)); then
		continue
	    fi
	    # compute parameters
	    bpt=$((rsz*1024/bs))
	    blocks=$((size*((1024*1024)/bs)/crg))
	    count=$blocks
	    # show computed parameters
	    actual_rsz=$((bpt*bs/1024))
	    actual_size=$((bs*count*crg/1024))
	    str=`printf 'total_size %8dK rsz %4d crg %5d thr %5d ' \
		         $((actual_size*ndevs)) $actual_rsz $((crg*ndevs)) $((thr*ndevs))`
	    echo "==============> $str" >> $workf
	    print_summary -n "$str"
	    freemem=`awk < /proc/meminfo '/^MemTotal:/ {printf "%d\n", $2}'`
	    if (((actual_rsz*thr/crg + 64)*crg*ndevs > freemem)); then
		print_summary "ENOMEM"
		continue
	    fi
	    # run tests
	    for action in $actions; do
		print_summary -n "$action "
		echo "=====> $action" >> $workf
		tmpf=${workf}_tmp
                # start test
		t0=`date +%s.%N`
		for ((i=0;i<ndevs;i++)); do
		    dev=${devs[i]}
		    devsize=$((bs*`sg_readcap -b ${dev} | awk '{print $1}'`/1024))
		    if [ $devsize -lt $actual_size ]; then
			_dev=$(sg_map | grep $dev | awk '{ print $2; }')
			echo -e "device $_dev not big enough: $devsize <" \
				"$actual_size.\nConsider reducing \$size"
			exit 1
		    fi
		    if [ $action = read ]; then
			inf="if=$dev"
			outf="of=/dev/null"
			skip=skip
		    else
			inf="if=/dev/zero"
			outf="of=$dev"
			skip=seek
		    fi
		    for ((j=0;j<crg;j++)); do 
			sgp_dd 2> ${tmpf}_${i}_${j} \
			    $inf $outf ${skip}=$((1024+j*blocks)) \
			    thr=$((thr/crg)) count=$count bs=$bs bpt=$bpt time=1&
		    done
		done 
		wait
		t1=`date +%s.%N`
	        # collect/check individual stats
		echo > $tmpf
		ok=0
		for ((i=0;i<ndevs;i++)); do
		    for ((j=0;j<crg;j++)); do
			rtmp=${tmpf}_${i}_${j}
			if grep 'error' $rtmp > /dev/null 2>&1; then
				echo "Error found in $rtmp"
			elif grep 'time to transfer data' $rtmp > /dev/null 2>&1; then
			    ok=$((ok + 1))
			fi
			cat ${rtmp} >> $tmpf
			cat ${rtmp} >> $workf
			rm  ${rtmp}
		    done
		done
		if ((ok != ndevs*crg)); then
		    print_summary -n "$((ndevs*crg - ok)) failed "
		else
	            # compute MB/sec from elapsed
		    bw=`awk "BEGIN {printf \"%7.2f MB/s\", $actual_size * $ndevs / (( $t1 - $t0 ) * 1024); exit}"`
	            # compute MB/sec from nregions*slowest
		    check=`awk < $tmpf \
			'/time to transfer data/ {mb=$8/1.048576; if (n == 0 || mb < min) min = mb; n++}\
			END {printf "%5d x %6.2f = %7.2f MB/s", n, min, min * n}'`
		    print_summary -n "$bw $check "
		fi
		rm $tmpf
	    done
	    print_summary ""
	done
    done
done

if $sg_was_loaded; then
    echo "unloading sg module"
    rmmod sg
fi