#!/bin/bash 
# iotest v0.2 - lawalsh(at)tlinx.org : open usage allowed
# (c) 2014-2018
# vim=:SetNumberAndWidth

_prgpth="${0:?}"; _prg="${_prgpth##*/}"; _prgdr="${_prgpth%/$_prg}"
[[ -z $_prgdr || $_prg == $_prgdr ]] && $_prgdr="$PWD"
export PATH="$_prgdr:$_prgdr/lib:$PATH"
shopt -s expand_aliases extglob sourcepath ; set -o pipefail

#include stdalias
#include Types

Dd=$(type -P dd)

[[ $Dd ]] || { echo "Cannot find dd.  Cannot proceed.";  exit 1; }

alias my=declare sub=function int='my -i' array='my -a' 
alias map='my -A' intConst='int -x' string=my

# 1 num = block size
# num-num = range of block sizes to test; w/increment = "2x", so
# 4M-16M tests 4M, 8M, 16M
# 4M-12M test 4M, 8M, 12M
# count adjusted to xfer 4G, rounding up
#----

#all xfers are using 'devices' (/dev/zero for source, /dev/null for target)
# remote filenames "zero" and "null" should be setup to be remote devices

intConst K=1024 M=K*K G=M*K T=G*K
my sfxs="KMGT"
my prog='/bin/dd'

##
# Defaults for BS count and IOSIZE
##
check_parm () {
  my param=$1
  if [[ ${!param} =~ ^[0-9]+[$sfxs]$ ]]; then
    my unit=${!param:0-1}
    eval $param="${!param:0:0-1}*$unit"
  elif [[ ${!param} =~ ^[0-9]+$ ]]; then
    : # ok as is
  elif [[ ! ${!param} =~ ^[0-9]+\*[$sfxs]$ ]]; then
    printf >&2 "Invalid number format \"%s\", for \"%s\"\n" "${!param}" "$param"
    exit 1
  fi
}

my BS COUNT IOSIZE
for p in BS COUNT IOSIZE; do
  if [[ ${!p} ]]; then
   check_parm $p
   int $p=${!p}
 fi
done

int params=0
array -i coef=(0 0 0)
#: ${BS:=16*M}  ${COUNT:=64}
if ((${#BS})); then int BS=$BS; params+=1; coef[0]=1; fi
if ((${#COUNT})); then int COUNT=$COUNT; params+=1; coef[1]=1; fi
if ((${#IOSIZE})); then int IOSIZE=$IOSIZE; params+=1 coef[2]=1;  fi
int BS=${BS:-16*M}
int COUNT=${COUNT:-64}
int IOSIZE=${IOSIZE:-2*G}
if ((params>=3)); then 
  printf >&2 "Error: can only specify 2 params; 3rd is calculated\n"; exit 1;
elif ((params==2)); then
  if ((!coef[2])); then int IOSIZE=COUNT*BS
  elif ((!coef[1])); then int COUNT=IOSIZE/BS
  elif ((!coef[0])); then int BS=IOSIZE/COUNT
  fi
elif ((params==1)); then
  if ((coef[2])) ; then   # only size spec'd, use default for BS unless too big
    if ((BS>IOSIZE)); then BS=IOSIZE/COUNT
    else COUNT=IOSIZE/BS
    fi
  elif ((coef[1])); then BS=IOSIZE/COUNT
  else COUNT=IOSIZE/BS
  fi
fi

int IOSIZE=COUNT*BS
#echo "BS=$BS, COUNT=$COUNT, IOSIZE=$IOSIZE"

# desuffix  arg1 [arg2]
# desuffix  arg1 - num + suffix -> convert to int
#           arg2 - optional buff name (else print to stdout)
#           return 0 if no error
desuffix () {            #convert num+Suff => int store in optional Buff
  my str="${1:?}" ; shift
  my bufnam=""; (($#)) && bufnam=$1
  if [[ $str =~ ^([0-9]+)([KMGT])$ ]]; then 
    int num=${BASH_REMATCH[1]}*${BASH_REMATCH[2]}
    ((num)) || return 1
    if [[ $bufnam ]] ; then printf -v $bufnam "%d" "$num"
    else printf "%d" "$num" ; fi
  else 
    return $p 
  fi
}


hdisp () {
  my parm=${1:?}
  if [[ ! $parm =~ ^[0-9]+$ ]]; then echo >&2 "hdisp: parm1 is not num: \"$parm\""; fi
  int num=${1:?}; shift
  export PERL5OPT="$PERL5OPT -I/h/bin/lib"
  my buf="$(perl -e 'use Hout; use P; P "%s", human($ARGV[0]);' "$num")"
  #printf "hdisp-buf=%s\n" "$buf"
  string bufnam=""; (($#)) && bufnam=$1
  if [[ $bufnam ]] ; then printf -v $bufnam "%s" "$buf"
  else printf "%s" "$buf" ; fi
}

map args=([b]=setblocksize [i]=iosize )

setblocksize () { BS=$1 ; }

iosize () { IOSIZE=$1 ; }

check_params () {
  int num=0
  while (($#)) ; do
    my arg=$1 ; shift
    if [[ ${arg:0:1} == - ]]; then arg=${arg:1}; fi
    my switch=${args["$arg"]:-""}
    int iswitch=0
    if [[ ${switch:-""} ]]; then
      my val=$1; shift
      my nval=""
      int ival=0
      desuffix "$val" nval
      if [[ $nval ]]; then ival=0+nval; fi
      $switch $ival
    fi
  done
  COUNT=IOSIZE/BS
}

(($#)) && check_params "$@"

string testdir=/h

array reada=($testdir/zero /dev/null )
array writea=(/dev/zero $testdir/null oflag=direct conv=nocreat,notrunc)

sub dd_io  {
  local if="$1" of="$2"; shift 2
  nice --19 $Dd if="$if" of="$of" bs="$BS" count="$COUNT" iflag=fullblock\
      conv=nocreat "$@"
}

dd () {
  local if="$1" of="$2" ; parms="$3"; shift 2
  array out=()
  my ignore_RE='^[^ \t]+ records\ (in|out)$'
  readarray out < <(  dd_io "$if" "$of" "$@" |& { int s=$?
                        if ((s)); then echo >&2 "stat:$s"; else cat; fi ; }
                    )
  printf "%s\n" "${out[@]}"
  return 0
}
  
dd_format () { 
  my bytes btxt pnum1 suffp1 pnum2 suffp2 copt time rest
  while read bytes btxt pnum1 suffp1 pnum2 suffp2 copt time rest; do
    if ! [[ $time =~ ^[0-9] ]]; then
      copt=$pnum2; time=$suffp2;
    fi
    [[ $btxt == records ]] && continue
    [[ $bytes && $time ]] || continue
    my sizeb="" rateb=""
    my bps="$(echo -E "$bytes/$time" | bc)"
    hdisp "$bytes" sizeb
    hdisp "$bps" rateb
    my fmt="%d bytes (%sB) copied, %s s, %sB/s\n"
    printf "$fmt" "$bytes" "$sizeb" "$time" "$rateb"
  done 
}

onecycle () {
  echo -n "R:"; { dd "${reada[@]}" || exit $?; } | dd_format
  sync;sleep .1;sync
  echo -n "W:"; { dd "${writea[@]}" || exit $?; } | dd_format 
}

my bsbuf="" ios_buf=""

hdisp "$BS" bsbuf
hdisp "$IOSIZE" ios_buf

#if (($#)) ; then
# if [[ $1 == -h || $1 == -\? ]]; then
#   printf "Defaults: BS=16M, COUNT=%d, IOSIZE=%s" $Count $IOSize
#   printf "BlkSize, cnt & Total can be overriden with BS, COUNT & IOSIZE:\n"
#   printf "BS=16M (uses 1G as Total IOSIZE), or\n" "$_prg" :
#   printf "BS=1M COUNT=1K %s (uses 1G as Total IOSIZE), or\n" "$_prg" 
#   printf ""
# fi
#fi

printf "Using bs=%s, count=%s, iosize=%s\n" "$bsbuf" "$COUNT" "$ios_buf"


onecycle