#!/bin/bash

set -eu

COUNT_DISK_REWRITES="${COUNT_DISK_REWRITES:-1}"

command -v fio >/dev/null 2>&1 || { echo >&2 "fio required but it's not installed.  Aborting."; exit 1; }
command -v iostat >/dev/null 2>&1 || { echo >&2 "iostat required but it's not installed.  Aborting."; exit 1; }
command -v fdisk >/dev/null 2>&1 || { echo >&2 "fdisk required but it's not installed.  Aborting."; exit 1; }
command -v python >/dev/null 2>&1 || { echo >&2 "python required but it's not installed.  Aborting."; exit 1; }
command -v lsblk >/dev/null 2>&1 || { echo >&2 "lsblk required but it's not installed.  Aborting."; exit 1; }
command -v bc >/dev/null 2>&1 || { echo >&2 "bc required but it's not installed.  Aborting."; exit 1; }
command -v screen >/dev/null 2>&1 || { echo >&2 "screen required but it's not installed.  Aborting."; exit 1; }
command -v tee >/dev/null 2>&1 || { echo >&2 "tee required but it's not installed.  Aborting."; exit 1; }


if [ $# -eq 0 ]; then
  lsblk
  echo

  DEVICES=$(lsblk -J | python -c 'import json,sys; j=json.load(sys.stdin); a=[item["name"] for item in j["blockdevices"] if item["mountpoint"] == None and ("children" not in item or True not in map(lambda x: x["mountpoint"] != None, item["children"])) ]; print " ".join(a)')

else
  DEVICES="$@"
fi

echo
echo "Starting test on devices:" $DEVICES
read -r -p "Are you sure? [y/N] " response
case "$response" in
    [yY][eE][sS]|[yY])
        echo
        ;;
    *)
        exit
        ;;
esac

rm -f ./fio_bench.sh
cat <<\EOF >> ./fio_bench.sh
#!/bin/bash

DISK=$1
DISK_NAME=$(basename ${DISK})

COUNT=$2

SIZE_GB=$(sudo fdisk -l ${DISK} | grep 'Disk /dev/' | awk '{ print $3; }')
SIZE_MB=$(echo "1024 * ${SIZE_GB}" | bc)

#echo deadline | sudo tee /sys/block/${DISK_NAME}/queue/scheduler
QUEUE_DEPTH=$(cat /sys/block/${DISK_NAME}/queue/nr_requests)

while [ $COUNT -gt 0 ]; do
  echo "#### $(date '+%s')" | tee -a ./fio_${DISK_NAME}.txt
  sudo fio --name=test \
           --filename="${DISK}" \
           --ioengine=libaio \
           --bs=4M \
           --iodepth=${QUEUE_DEPTH} \
           --direct=1 \
           --rw=rw \
           --rwmixwrite=50 \
        | tee -a ./fio_${DISK_NAME}.txt 2>&1
  let COUNT-=1
done
EOF

screen -S iostat -d -m bash -c 'iostat -xydt 60 | tee -a ./iostat.txt 2>&1'

for var in $DEVICES; do
    echo "Starting test in screen on ${var}"
    screen -S fio_${var} -d -m bash ./fio_bench.sh /dev/${var} ${COUNT_DISK_REWRITES}
done