#!/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="$2"

# 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!";
    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
        echo ""
        if ! command -v fio &> /dev/null; then
            dd_speedtest
        else
            fio_speedtest
        fi
    fi
}

function fio_write_single_random_4k {
    # Single 4k Random Writes

    # This is a single process doing random 4K writes. This is where the pain
    # really, really lives; it's basically the worst possible thing you can ask a
    # disk to do. Where this happens most frequently in real life: copying home
    # directories and dotfiles, manipulating email stuff, some database operations,
    # source code trees.

    # When I ran this test against the high-performance SSDs in my Ubuntu
    # workstation, they pushed 127MiB/sec. The server just beneath it in the rack
    # only managed 33MiB/sec on its "high-performance" 7200RPM rust disks... but
    # even then, the vast majority of that speed is because the data is being
    # written asynchronously, allowing the operating system to batch it up into
    # larger, more efficient write operations.

    # If we add the argument --fsync=1, forcing the operating system to perform
    # synchronous writes (calling fsync after each block of data is written) the
    # picture gets much more grim: 2.6MiB/sec on the high-performance SSDs but
    # only 184KiB/sec on the "high-performance" rust. The SSDs were about four
    # times faster than the rust when data was written asynchronously but a
    # whopping fourteen times faster when 

    # --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 working directory you're currently in.    
    # --ioengine=posixaio sets the mode fio interacts with the filesystem. POSIX is a standard Windows, Macs, Linux, and BSD all understand, so it's great for portability—although inside fio itself, Windows users need to invoke --ioengine=windowsaio, not --ioengine=posixaio, unfortunately. AIO stands for Asynchronous Input Output and means that we can queue up multiple operations to be completed in whatever order the OS decides to complete them. (In this particular example, later arguments effectively nullify this.)
    # --rw=randwrite means exactly what it looks like it means: we're going to do random write operations to our test files in the current working directory. Other options include seqread, seqwrite, randread, and randrw, all of which should hopefully be fairly self-explanatory.
    # --bs=4k blocksize 4K. These are very small individual operations. This is where the pain lives; it's hard on the disk, and it also means a ton of extra overhead in the SATA, USB, SAS, SMB, or whatever other command channel lies between us and the disks, since a separate operation has to be commanded for each 4K of data.
    # --size=1g our test file(s) will be 1GB in size apiece. (We're only creating one, see next argument.)
    # --numjobs=1 we're only creating a single file, and running a single process commanding operations within that file. If we wanted to simulate multiple parallel processes, we'd do, eg, --numjobs=16, which would create 16 separate test files of --size size, and 16 separate processes operating on them at the same time.
    # --iodepth=1 this is how deep we're willing to try to stack commands in the OS's queue. Since we set this to 1, this is effectively pretty much the same thing as the sync IO engine—we're only asking for a single operation at a time, and the OS has to acknowledge receipt of every operation we ask for before we can ask for another. (It does not have to satisfy the request itself before we ask it to do more operations, it just has to acknowledge that we actually asked for it.)
    # --runtime=15 --time_based Run and even if we complete sooner, just start over again and keep going until 60 seconds is up.
    # --end_fsync=1 After all operations have been queued, keep the timer going until the OS reports that the very last one of them has been successfully completed—ie, actually written to disk.
    echo ""
    echo "Single 4K Random Writes (size=1G, time=15sec, jobs=1, iodepth=1)"
    x=`sudo fio \
        --name=fio-write-random-4k \
        --directory=$path \
        --ioengine=posixaio \
        --rw=randwrite \
        --bs=4k \
        --size=1g \
        --numjobs=1 \
        --iodepth=1 \
        --time_based --runtime=15 \
        --end_fsync=1`
    echo "  - $x " | /usr/bin/grep -A1 'Run status group' | tail -n1
    
    # 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. For example, the 4K single process test we tried first
    # wrote a pretty consistent 11MiB/sec on my MacBook Air's internal drive—but
    # this 16-process job fluctuated between about 10MiB/sec and 300MiB/sec during
    # the run, finishing with an average of 126MiB/sec.

    # 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.
    echo ""
    echo "Parallel 64K Random Writes (size=1G, time=15sec, jobs=16, iodepth=16)"
    x=`sudo fio \
        --name=fio-write-random-64k \
        --directory=$path \
        --ioengine=posixaio \
        --rw=randwrite \
        --bs=64k \
        --size=64m \
        --numjobs=16 \
        --iodepth=16 \
        --time_based --runtime=15 \
        --end_fsync=1`
    echo "  - $x " | /usr/bin/grep -A1 'Run status group' | tail -n1

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

function fio_write_single_sequential_1m {
    # Single 1M Random 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.

    echo ""
    echo "Single 1M Sequential Writes (size=1G, time=15sec, jobs=1, iodepth=1)"
    x=`sudo fio \
        --name=fio-write-random-1m \
        --directory=$path \
        --ioengine=posixaio \
        --rw=write \
        --bs=1m \
        --size=1g \
        --numjobs=1 \
        --iodepth=1 \
        --time_based --runtime=15 \
        --end_fsync=1`
    echo "  - $x " | /usr/bin/grep -A1 'Run status group' | tail -n1

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

function fio_read_sequential_1m {
    # Sequential Parallel Reads

    echo ""
    echo "Sequential 4x 1M Reads"
    x=`sudo fio \
        --name=fio-read-sequential-1m \
        --directory=$path \
        --ioengine=posixaio \
        --bs=1M \
        --numjobs=4 \
        --size=256M \
        --time_based --runtime=30s \
        --ramp_time=2s \
        --direct=1 \
        --verify=0 \
        --iodepth=64 \
        --rw=read \
        --group_reporting=1 \
        --iodepth_batch_submit=64 \
        --iodepth_batch_complete_max=64`
    echo "  - $x " | /usr/bin/grep -A1 'Run status group' | tail -n1
    rm -rf $path/fio-read-sequential-1m*
}

function fio_read_random_4k {
    # Random 4k Reads

    echo ""
    echo "Random 4k Reads"
    x=`sudo fio \
        --name=fio-read-random-4k \
        --directory=$path \
        --ioengine=posixaio \
        --rw=randread \
        --bs=4k \
        --size=1g \
        --time_based --runtime=30s \
        --ramp_time=2s \
        --direct=1 \
        --verify=0 \
        --iodepth=256 \
        --rw=read \
        --group_reporting=1 \
        --iodepth_batch_submit=256 \
        --iodepth_batch_complete_max=256`
    echo "  - $x " | /usr/bin/grep -A1 'Run status group' | tail -n1
    rm -rf $path/fio-read-random-4k*
}

function fio_speedtest {
    # 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"
    exit 0
}

# Go
main
