#!/bin/bash

# Script name: fio_throughput_test.sh

network_mount="/data/user_data/io_test"
log_dir="$HOME/log/fio-tests"
hostname=$(hostname)
date=$(date +%Y-%m-%d)
log_file="${log_dir}/fio-${hostname}-${date}.log"

# Default values
file_size="8G"
iodepth="16" # Number of I/O operations to keep in flight. Higher values can increase performance but may also increase latency.
numjobs="4" # Adjust based on the number of parallel streams you wish to test
runtime="300" # Optional: limits each test to 5 minutes
ioengine="libaio" # Available ioengine types (sync, psync, vsync, libaio, posixaio, mmap, rbd, net, null)

# Function to display help message
show_help() {
    echo "Usage: fio_throughput_test.sh [-s file_size] [-d iodepth] [-j numjobs] [-r runtime] [-e ioengine] [-m network_mount] [-h]"
    echo
    echo "Options:"
    echo "  -s file_size      Size of the file to test with (default: 8G)"
    echo "  -d iodepth        Number of I/O operations to keep in flight (default: 16)"
    echo "  -j numjobs        Number of parallel streams to test (default: 4)"
    echo "  -r runtime        Runtime for each test in seconds (default: 300)"
    echo "  -e ioengine       I/O engine to use (default: libaio)"
    echo "  -m network_mount  Network mount point (default: /data/user_data/io_test)"
    echo "  -h                Show this help message"
    exit 0
}

# Parse command-line arguments
while getopts "s:d:j:r:e:m:h" flag; do
    case "${flag}" in
        s) file_size=${OPTARG};;
        d) iodepth=${OPTARG};;
        j) numjobs=${OPTARG};;
        r) runtime=${OPTARG};;
        e) ioengine=${OPTARG};;
        m) network_mount=${OPTARG};;
        h) show_help;;
        *) show_help;;
    esac
done

mkdir -p ${log_dir}

# Function to run fio test
run_fio_test() {
    local test_name=$1
    local rw_pattern=$2
    local block_size=$3
    local filename="${network_mount}/fio_${test_name}_testfile"

    echo "Running fio ${test_name} test..."
    fio_output=$(fio --name=${test_name} --ioengine=${ioengine} --iodepth=${iodepth} --rw=${rw_pattern} --bs=${block_size} --direct=1 --size=${file_size} --numjobs=${numjobs} --runtime=${runtime} --filename=${filename} --group_reporting 2>&1 | tee -a ${log_file})
    throughput=$(echo "$fio_output" | grep -E "READ:|WRITE:" | awk -F 'bw=' '{print $2}' | awk '{print $1$2}')

    if [[ $? -ne 0 ]]; then
        echo "Error: fio ${test_name} test failed." | tee -a ${log_file}
        return 1
    fi

    echo "${test_name} Throughput: ${throughput}" | tee -a ${log_file}
    echo "${test_name} Throughput: ${throughput}" >> throughput_summary.log
}

# Check if fio is installed
if ! command -v fio &> /dev/null; then
    echo "Error: fio is not installed." | tee -a ${log_file}
    exit 1
fi

# Run tests
run_fio_test "random_write" "randwrite" "4k"
run_fio_test "random_read" "randread" "4k"
run_fio_test "sequential_write" "write" "1M"
run_fio_test "sequential_read" "read" "1M"

# Summarize throughput
echo "Summary of Throughput:" | tee -a ${log_file}
cat throughput_summary.log | tee -a ${log_file}
rm throughput_summary.log