#!/usr/bin/env bash

# Robust HD/SDD/NVMe performance CLI utility
# Utilizing FIO for sequential/random writes/writes
# Dependencies: fio (apt install fio)
# See: https://cloud.google.com/compute/docs/disks/benchmarking-pd-performance
# See: https://arstechnica.com/gadgets/2020/02/how-fast-are-your-disks-find-out-the-open-source-way-with-fio/
# mReschke 2024-01-18

# CLI Parameters
path="$1"
option=""
simple=false
engine=""        # IO engine; empty = auto-detect (io_uring > libaio > posixaio > sync)
direct=""        # O_DIRECT; empty = auto-detect; 1 = bypass page cache, 0 = buffered
runtime=30       # seconds per test (longer = more accurate, exposes SLC cache exhaustion)
for arg in "${@:2}"; do
    case "$arg" in
        --simple) simple=true ;;
        --dd) option="--dd" ;;
        --fio) option="--fio" ;;
        --buffered) direct=0 ;;
        --direct) direct=1 ;;
        --engine=*) engine="${arg#*=}" ;;
        --direct=*) direct="${arg#*=}" ;;
        --runtime=*) runtime="${arg#*=}" ;;
    esac
done

# Main application flow
function main {

    # Show usage if no params
    if [ ! "$path" ]; then
        usage
    fi

    # Understand . path
    if [ "$path" == '.' ]; then
        path=$(pwd)
    fi

    # Check if path exists
    if [ ! -e "$path" ]; then
        echo "Path $path does not exist"
        exit 1
    fi

    # Must type y or n THEN press enter (which I like better)
    echo "NOTICE: 1GB free space on '$path' is required to perform the benchmark."
    echo -n "Are you ready to start a robust IO benchmark against '$path' ?"; read answer
    if [ "$answer" != "${answer#[Yy]}" ]; then
        echo "Great! Starting benchmark now!"
        echo "------------------------------"
    else
        echo "Ok, cancelled!"
        exit 0
    fi

    # Use dd of fio based on param or defaults
    if [ "$option" == "--dd" ]; then
        dd_speedtest
    elif [ "$option" == "--fio" ]; then
        fio_speedtest
    elif [ "$option" == "" ]; then
        # If fio is installed, use it, else use dd
        [ "$simple" != true ] && echo ""
        if ! command -v fio &> /dev/null; then
            dd_speedtest
        else
            fio_speedtest
        fi
    fi
}

function print_result {
    # Print a single fio test result, either as the full "Run status group"
    # line or, when --simple is given, as a compact aligned "Label  value" row.
    # The simple value is taken from the parenthetical (MB/s) figure that
    # follows the first bw= value in fio's output.
    local label="$1"
    local output="$2"
    local status
    status=$(echo "  - $output " | /usr/bin/grep -A1 'Run status group' | tail -n1)
    if [ "$simple" == true ]; then
        local value
        # Grab the parenthetical (MB/s) figure after the first bw=, then put a
        # space between the number and the unit, eg "98.4MB/s" -> "98.4 MB/s".
        value=$(echo "$status" | /usr/bin/grep -oP 'bw=\S+\s+\(\K[^)]+' | sed -E 's/([0-9.]+)([A-Za-z])/\1 \2/')
        printf "%-28s%s\n" "$label" "$value"
    else
        echo "$status"
    fi
}

function fio_write_single_random_4k {
    # 4K Random Writes (parallel, deep queue)

    # Random 4K writes are the worst thing you can ask a disk to do. Where this
    # happens most in real life: copying home directories and dotfiles, email
    # stores, some database operations, source code trees.

    # IMPORTANT: the original version of this test used numjobs=1 + iodepth=1.
    # That measures single-op LATENCY (the absolute worst case): the OS issues
    # one 4K write, waits for it to be acknowledged, then issues the next. On
    # NVMe that lands around ~3000 IOPS / ~12MB/sec -- which looks alarmingly
    # slow but is NOT a bug. NVMe's entire advantage is deep command queues and
    # parallelism, so a queue depth of 1 deliberately measures the one scenario
    # NVMe is bad at. To measure the drive's actual capability we now drive it
    # with multiple jobs and a deep queue (4 jobs x iodepth=64), keeping
    # hundreds of operations in flight at once. Expect this to jump into the
    # hundreds of MB/sec (often 1GB/sec+) on a healthy NVMe pool.

    # --name= is a required argument, but it's basically human-friendly fluff—fio will create files based on that name to test with, inside the directory you point it at.
    # --directory= where the test files are created; this is the path/mount you are benchmarking.
    # --ioengine=$engine the engine fio uses to talk to the kernel, auto-detected by detect_io_settings(). On modern Linux this is io_uring (lowest overhead); it falls back to libaio, then posixaio, then sync. NOTE: libaio is only truly asynchronous with --direct=1; with buffered IO it silently degrades to synchronous, which is exactly why io_uring is preferred.
    # --rw=randwrite random write operations. Other options include read, write (sequential), randread, and randrw.
    # --bs=4k blocksize 4K. These are very small individual operations. This is where the pain lives; it's hard on the disk and adds a ton of command-channel overhead, since a separate operation has to be commanded for each 4K of data.
    # --size=256m each job's test file is 256MB. With --numjobs=4 that's 4 files totaling ~1GB.
    # --numjobs=4 run 4 parallel processes, each with its own test file, to feed the drive from multiple threads at once (single-threaded cannot saturate NVMe).
    # --iodepth=64 how many commands each job keeps stacked in the queue at once. Deep queues are how NVMe reaches its real IOPS; 4 jobs x 64 = up to 256 ops in flight.
    # --direct=$direct O_DIRECT, auto-detected. 1 bypasses the OS page cache so we measure the device instead of RAM; 0 (buffered) is the fallback for filesystems/mounts that reject O_DIRECT (eg older OpenZFS, some NFS).
    # --ramp_time=2s discard the first 2 seconds of results so we measure steady state, not the initial warm-up burst.
    # --runtime=$runtime --time_based run for this many seconds, looping over the file(s) if we finish early.
    # --group_reporting=1 aggregate all jobs into a single combined result line instead of one line per job.
    # --end_fsync=1 after the timed run, flush everything to stable storage and count that time, so cached writes can't inflate the number.
    [ "$simple" != true ] && { echo ""; echo "4K Random Writes (bs=4k, jobs=4, iodepth=64, engine=$engine, direct=$direct, ${runtime}s)"; }
    x=`sudo fio \
        --name=fio-write-random-4k \
        --directory=$path \
        --ioengine=$engine \
        --rw=randwrite \
        --bs=4k \
        --size=256m \
        --numjobs=4 \
        --iodepth=64 \
        --direct=$direct \
        --ramp_time=2s \
        --time_based --runtime=$runtime \
        --group_reporting=1 \
        --end_fsync=1`
    print_result "4K Random Writes" "$x"

    # Cleanup my test files
    rm -rf $path/fio-write-random-4k*
}

function fio_write_parallel_random_64k {
    # Parallel 64k Random Writes

    # This time, we're creating 16 separate 64MB files (still totaling 1GB, when
    # all put together) and we're issuing 64KB blocksized random write operations.
    # We're doing it with sixteen separate processes running in parallel, and
    # we're queuing up to 16 simultaneous asynchronous ops before we pause and wait
    # for the OS to start acknowledging their receipt.

    # This is a pretty decent approximation of a significantly busy system. It's
    # not doing any one particularly nasty thing—like running a database engine or
    # copying tons of dotfiles from a user's home directory—but it is coping with
    # a bunch of applications doing moderately demanding stuff all at once.

    # This is also a pretty good, slightly pessimistic approximation of a busy,
    # multi-user system like a NAS, which needs to handle multiple 1MB operations
    # simultaneously for different users. If several people or processes are trying
    # to read or write big files (photos, movies, whatever) at once, the OS tries
    # to feed them all data simultaneously. This pretty quickly devolves down to a
    # pattern of multiple random small block access. So in addition to "busy desktop
    # with lots of apps," think "busy fileserver with several people actively using it."

    # You will see a lot more variation in speed as you watch this operation play
    # out on the console. Unlike the steady trickle you'd get from a queue-depth-1
    # single-op latency test, this 16-process job can fluctuate wildly—eg between
    # about 10MiB/sec and 300MiB/sec during the run—as the OS and SSD firmware
    # catch good and bad luck aggregating writes.

    # Most of the variation you're seeing here is due to the operating system and
    # SSD firmware sometimes being able to aggregate multiple writes. When it
    # manages to aggregate them helpfully, it can write them in a way that allows
    # parallel writes to all the individual physical media stripes inside the SSD.
    # Sometimes, it still ends up having to give up and write to only a single
    # physical media stripe at a time—or a garbage collection or other maintenance
    # operation at the SSD firmware level needs to run briefly in the background,
    # slowing things down.
    [ "$simple" != true ] && { echo ""; echo "Parallel 64K Random Writes (bs=64k, jobs=16, iodepth=16, engine=$engine, direct=$direct, ${runtime}s)"; }
    x=`sudo fio \
        --name=fio-write-random-64k \
        --directory=$path \
        --ioengine=$engine \
        --rw=randwrite \
        --bs=64k \
        --size=64m \
        --numjobs=16 \
        --iodepth=16 \
        --direct=$direct \
        --ramp_time=2s \
        --time_based --runtime=$runtime \
        --group_reporting=1 \
        --end_fsync=1`
    print_result "Parallel 64K Random Writes" "$x"

    # Cleanup my test files
    rm -rf $path/fio-write-random-64k*
}

function fio_write_single_sequential_1m {
    # Single 1M Sequential Writes

    # This is pretty close to the best-case scenario for a real-world system
    # doing real-world things. No, it's not quite as fast as a single, truly
    # contiguous write... but the 1MiB blocksize is large enough that it's quite
    # close. Besides, if literally any other disk activity is requested simultaneously
    # with a contiguous write, the "contiguous" write devolves to this level of
    # performance pretty much instantly, so this is a much more realistic test of
    # the upper end of storage performance on a typical system.

    # You'll see some kooky fluctuations on SSDs when doing this test. This is largely
    # due to the SSD's firmware having better luck or worse luck at any given time,
    # when it's trying to queue operations so that it can write across all physical
    # media stripes cleanly at once. Rust disks will tend to provide a much more
    # consistent, though typically lower, throughput across the run.

    # You can also see SSD performance fall off a cliff here if you exhaust an
    # onboard write cache—TLC and QLC drives tend to have small write cache areas
    # made of much faster MLC or SLC media. Once those get exhausted, the disk has
    # to drop to writing directly to the much slower TLC/QLC media where the data
    # eventually lands. This is the major difference between, for example, Samsung
    # EVO and Pro SSDs—the EVOs have slow TLC media with a fast MLC cache, where
    # the Pros use the higher-performance, higher-longevity MLC media throughout
    # the entire SSD.

    # If you have any doubt at all about a TLC or QLC disk's ability to sustain
    # heavy writes, you may want to experimentally extend your time duration here.
    # If you watch the throughput live as the job progresses, you'll see the impact
    # immediately when you run out of cache—what had been a fairly steady,
    # several-hundred-MiB/sec throughput will suddenly plummet to half the speed
    # or less and get considerably less stable as well.

    # However, you might choose to take the opposite position—you might not
    # expect to do sustained heavy writes very frequently, in which case you
    # actually are more interested in the on-cache behavior. What's important
    # here is that you understand both what you want to test, and how to test
    # it accurately.

    [ "$simple" != true ] && { echo ""; echo "Single 1M Sequential Writes (bs=1m, jobs=1, iodepth=16, engine=$engine, direct=$direct, ${runtime}s)"; }
    x=`sudo fio \
        --name=fio-write-seq-1m \
        --directory=$path \
        --ioengine=$engine \
        --rw=write \
        --bs=1m \
        --size=1g \
        --numjobs=1 \
        --iodepth=16 \
        --direct=$direct \
        --ramp_time=2s \
        --time_based --runtime=$runtime \
        --group_reporting=1 \
        --end_fsync=1`
    print_result "Single 1M Sequential Writes" "$x"

    # Cleanup my test files
    rm -rf $path/fio-write-seq-1m*
}

function fio_read_sequential_1m {
    # Sequential Parallel Reads

    [ "$simple" != true ] && { echo ""; echo "Sequential 4x 1M Reads (bs=1m, jobs=4, iodepth=64, engine=$engine, direct=$direct, ${runtime}s)"; }
    x=`sudo fio \
        --name=fio-read-sequential-1m \
        --directory=$path \
        --ioengine=$engine \
        --bs=1M \
        --numjobs=4 \
        --size=256M \
        --time_based --runtime=$runtime \
        --ramp_time=2s \
        --direct=$direct \
        --verify=0 \
        --iodepth=64 \
        --rw=read \
        --group_reporting=1 \
        --iodepth_batch_submit=64 \
        --iodepth_batch_complete_max=64`
    print_result "Sequential 4x 1M Reads" "$x"
    rm -rf $path/fio-read-sequential-1m*
}

function fio_read_random_4k {
    # Random 4k Reads (parallel, deep queue)

    # 4K random reads are the highest-IOPS thing a drive does, but a single
    # submitting thread goes CPU-bound long before an NVMe runs out of capacity,
    # so it under-reports peak IOPS. We spread the work over 4 jobs x iodepth=64
    # (256 ops in flight, like the 4K write test) to saturate fast drives, while
    # still reporting the honest (low) number on a spinning HDD. size=256m per
    # job keeps the total footprint at ~1GB across the 4 jobs.
    [ "$simple" != true ] && { echo ""; echo "Random 4k Reads (bs=4k, jobs=4, iodepth=64, engine=$engine, direct=$direct, ${runtime}s)"; }
    x=`sudo fio \
        --name=fio-read-random-4k \
        --directory=$path \
        --ioengine=$engine \
        --rw=randread \
        --bs=4k \
        --size=256m \
        --numjobs=4 \
        --time_based --runtime=$runtime \
        --ramp_time=2s \
        --direct=$direct \
        --verify=0 \
        --iodepth=64 \
        --group_reporting=1 \
        --iodepth_batch_submit=64 \
        --iodepth_batch_complete_max=64`
    print_result "Random 4k Reads" "$x"
    rm -rf $path/fio-read-random-4k*
}

function fio_probe {
    # Run a tiny throwaway fio job to see if a given (engine, direct) combo
    # actually works on this path/filesystem. Returns 0 on success. This is how
    # we stay accurate AND portable across ext4, xfs, btrfs, ZFS, NFS, etc.
    # without the caller needing to know what each filesystem supports.
    local eng="$1" dir="$2"
    sudo fio --name=probe \
        --directory="$path" \
        --ioengine="$eng" \
        --rw=write --bs=4k --size=1m \
        --direct="$dir" \
        --time_based --runtime=1 \
        > /dev/null 2>&1
    local rc=$?
    rm -rf "$path"/probe* 2>/dev/null
    return $rc
}

function detect_io_settings {
    # Pick the fastest available IO engine unless the user forced one with
    # --engine=. io_uring is the modern, lowest-overhead Linux engine; libaio
    # is older (and only truly async with O_DIRECT); posixaio/sync are the
    # portable fallbacks (eg some NFS mounts).
    if [ -z "$engine" ]; then
        for eng in io_uring libaio posixaio sync; do
            if fio_probe "$eng" 0; then engine="$eng"; break; fi
        done
        [ -z "$engine" ] && engine="sync"
    fi

    # Decide O_DIRECT unless forced with --direct/--buffered. O_DIRECT bypasses
    # the OS page cache so we measure the device instead of RAM. Not every
    # filesystem supports it (eg older OpenZFS < 2.3, some NFS mounts), so we
    # probe and fall back to buffered rather than erroring out.
    if [ -z "$direct" ]; then
        if fio_probe "$engine" 1; then direct=1; else direct=0; fi
    fi

    if [ "$simple" != true ]; then
        echo ""
        echo "IO engine : $engine"
        if [ "$direct" == 1 ]; then
            echo "O_DIRECT  : enabled (bypassing page cache -- measuring the device)"
        else
            echo "O_DIRECT  : DISABLED (buffered) -- this filesystem/mount rejected"
            echo "            O_DIRECT, so results (especially reads) may reflect the"
            echo "            RAM cache rather than the underlying device. Pass"
            echo "            --direct to force it if you believe it should work."
        fi
    fi
}

function fio_speedtest {
    # Auto-detect the best engine and whether O_DIRECT works on this path.
    detect_io_settings

    # Write tests
    fio_write_single_random_4k
    fio_write_parallel_random_64k
    fio_write_single_sequential_1m

    # Read Tests
    fio_read_sequential_1m
    fio_read_random_4k
}

function dd_speedtest {
    # Basic HD speed test using DD
    # mReschke 2017-07-11
    
    file=$path/bigfile
    size=1024

    echo "Running dd based HD/SSD/NVMe Benchmarks"
    echo "---------------------------------------"

    printf "Cached write speed...\n"
    dd if=/dev/zero of=$file bs=1M count=$size

    printf "\nUncached write speed...\n"
    dd if=/dev/zero of=$file bs=1M count=$size conv=fdatasync,notrunc

    printf "\nUncached read speed...\n"
    echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null
    dd if=$file of=/dev/null bs=1M count=$size

    printf "\nCached read speed...\n"
    dd if=$file of=/dev/null bs=1M count=$size

    rm $file
    printf "\nDone\n"
}

# Show help and usage information
function usage {
    echo "Robust Flexible Input/Output HD Speedtest"
    echo "  If FIO is installed, we use FIO for more detailed performance analysis."
    echo "  If FIO is not installed, we use basic DD analysis."
    echo "  You should apt install fio (pacman -S fio) for detailed analysis."
    echo "mReschke 2024-01-18"
    echo ""
    echo "NOTICE, this creates a 1GB file on the desired destination disk."
    echo "Please ensure you have write access with 1GB free space on destination."
    echo ""
    echo "Usage:"
    echo "  This will use FIO if installed, else DD"
    echo "    ./speedtest-hd /mnt/somedisk"
    echo "    ./speedtest-hd ."
    echo ""
    echo "  This will force FIO"
    echo "    ./speedtest-hd /mnt/somedisk --fio"
    echo "    ./speedtest-hd . --fio"
    echo ""
    echo "  This will force DD"
    echo "    ./speedtest-hd /mnt/somedisk --dd"
    echo "    ./speedtest-hd . --dd"
    echo ""
    echo "  Add --simple (FIO only) for a compact, aligned summary of MB/s values"
    echo "    ./speedtest-hd . --simple"
    echo "    ./speedtest-hd . --fio --simple"
    echo ""
    echo "  FIO tuning flags (all auto-detected by default, override as needed):"
    echo "    --engine=io_uring|libaio|posixaio|sync   IO engine (default: auto)"
    echo "    --direct        Force O_DIRECT (bypass page cache)"
    echo "    --buffered      Force buffered IO (eg if O_DIRECT is unsupported)"
    echo "    --runtime=SEC   Seconds per test (default: 30)"
    echo "  Examples:"
    echo "    ./speedtest-hd /mnt/nvmepool --engine=io_uring --runtime=60"
    echo "    ./speedtest-hd /mnt/nfsshare --buffered"
    exit 0
}

# Go
main