#!/bin/bash

############################################################################
#
# Script for testing block device I/O performance. Running this script on a
# block device that is connected to a remote SCST target device allows to
# test the performance of the transport protocols implemented in SCST. The
# operation of this script is similar to iozone, while this script is easier
# to use.
#
# Copyright (C) 2009 Bart Van Assche <bart.vanassche@gmail.com>.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, version 2
# of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
############################################################################

#########################
# Function definitions  #
#########################

function usage {
  echo "Usage: $0 [-a] [-d] [-i <i>] [-n] [-r] [-s <l2s>] <dev>"
  echo "        -a - use asynchronous (buffered) I/O."
  echo "        -d - use direct (non-buffered) I/O."
  echo "        -i - number times each test is iterated."
  echo "        -n - do not verify the data on <dev> before overwriting it."
  echo "        -r - only perform the read test."
  echo "        -s - logarithm base two of the I/O size."
  echo "        <dev> - block device to run the I/O performance test on."
}

function drop_caches {
  sync
  if [ -w /proc/sys/vm/drop_caches ]; then
    echo 3 > /proc/sys/vm/drop_caches
  fi
}

# Read times in seconds from stdin, one number per line, echo each number
# using format $1, and also echo the average transfer size in MB/s, its
# standard deviation and the number of IOPS using the total I/O size $2 and
# the block transfer size $3.
function echo_and_calc_avg {
  awk -v fmt="$1" -v iosize="$2" -v blocksize="$3" '{if ($1 != 0){n++;sum+=iosize/$1;sumsq+=iosize*iosize/($1*$1)};printf fmt, $1} END{d=(n>0?sumsq/n-sum*sum/n/n:0);avg=(n>0?sum/n:0);stddev=(d>0?sqrt(d):0);iops=avg/blocksize;printf fmt fmt fmt,avg/2**20,stddev/2**20,iops}'
}

#########################
# Default settings      #
#########################

iterations=3
log2_io_size=30       # 1 GB
log2_min_blocksize=9  # 512 bytes
log2_max_blocksize=26 # 64 MB
iotype=direct
read_test_only=false
verify_device_data=true


#########################
# Argument processing   #
#########################

set -- $(/usr/bin/getopt "adi:nrs:" "$@")
while [ "$1" != "${1#-}" ]
do
  case "$1" in
    '-a') iotype="buffered"; shift;;
    '-d') iotype="direct"; shift;;
    '-i') iterations="$2"; shift; shift;;
    '-n') verify_device_data="false"; shift;;
    '-r') read_test_only="true"; shift;;
    '-s') log2_io_size="$2"; shift; shift;;
    '--') shift;;
    *)    usage; exit 1;;
  esac
done

if [ "$#" != 1 ]; then
  usage
  exit 1
fi

device="$1"


####################
# Performance test #
####################

if [ ! -e "${device}" ]; then
  echo "Error: device ${device} does not exist."
  exit 1
fi

if [ "${read_test_only}" = "false" -a ! -w "${device}" ]; then
  echo "Error: device ${device} is not writeable."
  exit 1
fi

if [ "${read_test_only}" = "false" -a "${verify_device_data}" = "true" ] \
   && ! cmp -s -n $((2**log2_io_size)) "${device}" /dev/zero
then
  echo "Error: device ${device} still contains data."
  exit 1
fi

if [ "${iotype}" = "direct" ]; then
  dd_oflags="oflag=direct"
  dd_iflags="iflag=direct"
else
  dd_oflags="oflag=sync"
  dd_iflags=""
fi

# Header
printf "%9s " blocksize
for ((i = 0; i < ${iterations}; i++))
do
  printf "%6s " "W(s)"
done
printf "%6s %6s %6s" "W(avg,MB/s)" "W(stddev,MB/s)" "W(IOPS)"
for ((i = 0; i < ${iterations}; i++))
do
  printf "%6s " "R(s)"
done
printf "%6s %6s %6s" "R(avg,MB/s)" "R(stddev,MB/s)" "R(IOPS)"
printf "\n"

# Measurements
for ((log2_blocksize = log2_max_blocksize;
      log2_blocksize >= log2_min_blocksize;
      log2_blocksize--))
do
  if [ $log2_blocksize -gt $log2_io_size ]; then
    continue
  fi
  iosize=$((2**log2_io_size))
  bs=$((2**log2_blocksize))
  count=$((2**(log2_io_size - log2_blocksize)))
  printf "%9d " ${bs}
  for ((i = 0; i < ${iterations}; i++))
  do
    if [ "${read_test_only}" = "false" ]; then
      drop_caches
      dd if=/dev/zero of="${device}" bs=${bs} count=${count} \
                    ${dd_oflags} 2>&1 \
                 | sed -n 's/.* \([0-9.]*\) s,.*/\1/p'
    else
      echo 0
    fi
  done | echo_and_calc_avg "%8g " ${iosize} ${bs}

  for ((i = 0; i < ${iterations}; i++))
  do
    drop_caches
    dd if="${device}" of=/dev/null bs=${bs} count=${count} \
                  ${dd_iflags} 2>&1 \
               | sed -n 's/.* \([0-9.]*\) s,.*/\1/p'
  done | echo_and_calc_avg "%8g " ${iosize} ${bs}
  printf "\n"
done