#!/usr/bin/env bash
# ______          __ _            _       
# | ___ \        / _| |          | |      
# | |_/ /__ _ __| |_| |_ ___  ___| |_ ___ 
# |  __/ _ \ '__|  _| __/ _ \/ __| __/ __|
# | | |  __/ |  | | | ||  __/\__ \ |_\__ \
# \_|  \___|_|  |_|  \__\___||___/\__|___/
#                                       
set +ex
set -o pipefail

# Ensure we have ample size on the tmpfs mount in order to pull down
# a large enough repo to conduct metadata IO at a proper scale.
DEFAULT_TMPFSSIZE=$((10**7))
DEFAULT_MOUNTPATH=/mnt/ramdisk
DEFAULT_CWDIRNAME=cw-test
DO_CLEANUP=0
DO_ONLY_METADATA_TEST=0
DO_ONLY_DATA_TEST=0
DO_EXPAND_OUTPUT=0

USER_MOUNTPATH=""
MOUNTPATH=${DEFAULT_MOUNTPATH}
TESTDIR=${DEFAULT_CWDIRNAME}

usage() {
  echo -e "usage: $0 -s [-t] [-c] [-d] [-h] [-m] [-f]"
  echo -e '\t' "-s: The volume mount that is slow or degraded."
  echo -e '\t' "-t: The tmpfs directory. This will be used solely for metadata test."
  echo -e '\t' "-c: Perform cleanup of any remnant file or directory stubs from past tests."
  echo -e '\t' "-d: Test directory name under which tests will be conducted. Default: 'cw-test'."
  echo -e '\t' "-m: Run only metadata test."
  echo -e '\t' "-f: Run only data IO+BW test."
  echo -e '\t' "-e: Expand the output to show details. Can be noisy in some cases, necessary in others."
  echo -e '\t' "-h: Print this usage."
}

while getopts ":s:t:d:cmfeh" options; do
  case "${options}" in
    s)
      USER_MOUNTPATH=${OPTARG}
      test -z "${USER_MOUNTPATH}" && usage && exit 1;
      ;;
    t)
      MOUNTPATH=${OPTARG}
      ;;
    c)
      DO_CLEANUP=1
      ;;
    d)
      TESTDIR=${OPTARG}
      ;;
    m)
      DO_ONLY_METADATA_TEST=1
      ;;
    f)
      DO_ONLY_DATA_TEST=1
      ;;
    e)
      DO_EXPAND_OUTPUT=1
      ;;
    h)
      usage
      exit 0
      ;;
    :)
      echo "-${OPTARG} requires an argument."
      usage
      exit 1
      ;;
    *)
      echo "Unexpected argument"
      usage
      exit 1
      ;;
  esac
done

shift $((OPTIND-1))

TESTDIRPATH="${MOUNTPATH}/${TESTDIR}"
USERDIRPATH="${USER_MOUNTPATH%/}/${TESTDIR}"

## Run cleanup after all tests complete.
# This should ensure we don't take up any additional space created by the test files on
# the users volume. Any test directories or files we create, should be prefixed with
# "test" substring in order to make sure they are cleaned up automatically. Any stubs
# not prefixed with "test" will not be cleaned up automatically just to be safe, and
# thus will need manual cleanup.
cleanup() {
    popd || true
    echo "Cleaning up..."
    find ${TESTDIRPATH} -type d -name "cwtest*" -exec rm -rf {} \; 2>/dev/null
    find ${TESTDIRPATH} -type f -name "cwtest*" -exec rm -f {} \; 2>/dev/null
    test -d ${TESTDIRPATH} && rmdir ${TESTDIRPATH}
    find ${USERDIRPATH} -type d -name "cwtest*" -exec rm -rf {} \; 2>/dev/null
    find ${USERDIRPATH} -type f -name "cwtest*" -exec rm -f {} \; 2>/dev/null
    test -d ${USERDIRPATH} && rmdir ${USERDIRPATH}
    echo "complete!"
}
trap cleanup EXIT SIGTERM SIGINT

# The cleanup ensures that any directories or files that are test
# stubs and created for test purposes are cleaned up once the test
# runs complete.
if [ "${DO_CLEANUP}" -eq 1 ]; then
    TESTDIRPATH="${MOUNTPATH}/${DEFAULT_CWDIRNAME}"
    USERDIRPATH="${USER_MOUNTPATH%/}/${DEFAULT_CWDIRNAME}"
    # This will call the trap, which will call the cleanup() function.
    exit 0
fi

if [ ! -d "${USER_MOUNTPATH}" ]; then
    echo "Error: User volume mount '${USER_MOUNTPATH}' does not exist."
    exit 1
fi

MOUNTPATH=${MOUNTPATH%/}
if [ ! -d "${MOUNTPATH}" ]; then
    echo "Error: tmpfs mountpoint '${MOUNTPATH}' does not exist."
    exit 1
fi

# Check that we are on a mount-path that is mounted as tmpfs
# and we have ample space to actually pull down our git repo
# in order to test for metadata ops.
tmpfs_size=$(df --output=avail --type tmpfs ${MOUNTPATH} | tail -1)
if [ "${tmpfs_size}" -lt "${DEFAULT_TMPFSSIZE}" ]; then
    echo "Error: tmpfs mountpoint '${MOUNTPATH}' needs to be sized more than 10GB"
    exit 1
fi

# Create both test dir paths in tmpfs and on the respective volume.
mkdir -p "${TESTDIRPATH}"
mkdir -p "${USERDIRPATH}"

pushd "${USERDIRPATH}"
if [ "$(pwd)" != "${USERDIRPATH}" ]; then
    echo "Error: Failed running tests inside test path. Please report to CoreWeave team."
    exit 1
fi

### Test functions
test_metadataio() {
    echo "Running test: 'metadata IO'"
    which git >/dev/null
    if [ $? -ne 0 ]; then
        echo "Error: git not found installed."
        exit 1
    fi   

    # We do not want to use upstream kernel or any other repo that has non-deterministic
    # patterns that we cannot run repeated tests across. This is a static repo that will
    # never or should never change.
    GITCMD="git clone --quiet"
    GITREPO="https://github.com/coreweave/testlinux.git"
    DWNPATH="cwtestlinux"
    TIMEO=10m

    if [ "${DO_EXPAND_OUTPUT}" -eq 1 ]; then
        GITCMD="git clone"
    fi

    timeout ${TIMEO} ${GITCMD} ${GITREPO} "${TESTDIRPATH}/${DWNPATH}"
    if [ $? -ne 0 ]; then
        echo "Error: initial git clone timed out. Please report to CoreWeave team."
        exit 1
    fi

    start_time=$(date +%s%3N)
    timeout ${TIMEO} ${GITCMD} "${TESTDIRPATH}/${DWNPATH}" "${USERDIRPATH}/${DWNPATH}"
    end_time=$(date +%s%3N)
    duration_ms=$((end_time - start_time))
    if [ $? -ne 0 ]; then
        echo "Error: git clone to volume timed out in ${duration_ms} ms. Please report to CoreWeave team."
        exit 1
    fi

    echo "RESULT ==> git-clone test: ${duration_ms} ms"
}

FIOCMD="fio -direct=1 -ioengine=sync -numjobs=12 -runtime=60 -time_based -group_reporting -name=cwtest -refill_buffers -randrepeat=0"

check_install_fio() {
  which fio >/dev/null
  if [ $? -ne 0 ]; then
      echo "fio not found, installing ..."
      sudo apt update && sudo apt install -y fio
      if [ $? -ne 0 ]; then
        echo "Error: unable to install fio. In order to run performance tests you need to have fio installed."
        exit 1
      fi
  fi
}

test_fileio() {
    echo "Running test: 'data IO'"
    check_install_fio
    test $? -ne 0 && exit 1

    BSSIZE="-bs=4k -size=2G"
    for io_op in "-rw=write" "-rw=read"; do
        CMDRUN="${FIOCMD} ${io_op} ${BSSIZE}"
        if [ "${DO_EXPAND_OUTPUT}" -eq 1 ]; then
            $CMDRUN
        else
            $CMDRUN | grep -w -e 'read:' -e 'write:'
        fi
    done
}

test_filebw() {
    echo "Running test: 'data BW'"
    check_install_fio
    test $? -ne 0 && exit 1

    BSSIZE="-bs=4M -size=20G"
    for io_op in "-rw=write" "-rw=read"; do
        CMDRUN="${FIOCMD} ${io_op} ${BSSIZE}"
        if [ "${DO_EXPAND_OUTPUT}" -eq 1 ]; then
            $CMDRUN
        else
            $CMDRUN | grep -w -e 'read:' -e 'write:'
        fi
    done
}

# All test functions that we want to run within the script will need to be put
# into this array in order to actually run. Commented tests or test functions
# that are not a part of this array will not run.
test_fns=()
test "${DO_ONLY_DATA_TEST}" -eq 1 && test_fns+=(test_fileio test_filebw)
test "${DO_ONLY_METADATA_TEST}" -eq 1 && test_fns+=(test_metadataio)

## Run all tests.
echo "Running tests..."
for tfunc in "${test_fns[@]}"; do
    "$tfunc"
done