kdbench

#!/usr/bin/env bash

# kdbench(1) is a utility to test storageclass performance easily 

# Copyright 2022 SysNet4Admin (Hoon Jo)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

[[ -n $DEBUG ]] && set -x
set -eou pipefail
IFS=$'\n\t'

SELF_CMD="$0"
SCRACHFILE=~/.kdbench/kdbench.scra
YAMLFILE=~/.kdbench/kdbench.yaml
GLOBAL_STORAGE_SIZE="100" # unit Gi 
GLOBAL_File_NAME="kdbench.yaml"

exit_oky(){
   exit 0
}
exit_err() {
   exit 1
}

usage(){
cat <<EOF
USAGE:
  kdbench                           : Select storageclass with fzf 
  kdbench <STORAGECLASS NAME>       : Select storageclass manually (w/o fzf)
  kdbench -q, --quick               : Quick run mode        (w/ fzf)
  kdbench -d, --direct              : Direct fio write mode (w/ fzf) 
  kdbench -a  <rwo,rox,rwx,rwop>    : Set access mode       (w/ fzf)
  kdbench -s  <storage size>        : Set stroage size      (w/ fzf)
  kdbench -f  <YAML FILE>           : Run as a file         (w/ fzf)
  kdbench reset                     : Reset all of resources 
  kdbench -h, --help                : Show this message
EOF
}

delete_pvc_after_job_done() {
  echo -n "Waiting for job finished by ttlSecondsAfterFinished..."
  while true; do 
    if [[ "$(kubectl get job -l fio=kdbench 2>&1)" = *"No resources"* ]]; then
      echo -e "\nkdbench finsihed and pvc will delete"
      kubectl delete pvc -l fio=kdbench 
      exit_oky
    else
      # print dot for waiting status 
      echo -n ".";sleep 1
    fi 
  done 
}

run_kdbench(){
  # deploy kdbench 
  if [[ "$#" -eq 1 ]]; then
    :
    if [[ "${1}" == 'file' ]]; then 
      kubectl apply -f "${GLOBAL_File_NAME}"
    else 
      echo -e "error: check Yaml file name"
      exit_err
    fi
  else 
    # mostly run this command
    kubectl apply -f "$YAMLFILE"
  fi 
  # wait message for dbench 
  echo -n "Waiting for kdbench's load (up to 5mins)..."
  local count=0
  while true; do
    PODSTATUS=$(kubectl get po -l job-name=kdbench 2>&1)
    if [[ "$PODSTATUS" = *"Creating"* ]] || [[ "$PODSTATUS" = *"Pending"* ]]; then
      echo -n ".";sleep 1; ((count=count+1))
      if [[ $count -gt 300 ]]; then 
        echo -e "\n\n********* Error: Events as below *********"
        $KUBECTL describe pvc kdbench-pv-claim | grep Message -A 50
        exit_err
      fi 
    elif [[ "$PODSTATUS" = *"Running"* ]]; then
      break 
    else
      echo "Pod is abnormal status"
    fi 
  done 
  # kdbench result print out on stdout 
  kubectl logs -f job/kdbench
}

get_storageclasses() {
  $KUBECTL get storageclasses -o=name | cut -c 29-
}

set_storageclass() {
  sed -i'' -e "s|your-storageclass|${1}|g" "$YAMLFILE"
}

choose_n_run_storageclasses_interactive() {
  local choice
  choice="$(FZF_DEFAULT_COMMAND="${SELF_CMD}" \
    fzf --ansi --no-preview || true)"
  if [[ -z "${choice}" ]]; then
    echo 2>&1 "error: you did not choose any of the options"
    exit_err
  else
      # create yaml file 
      create_yaml_file
      # change mode (optional)
      if [[ "$#" -eq 1 ]]; then
        if [[ "${1}" == 'ReadOnlyMany' || "${1}" == 'ReadWriteMany' || "${1}" == 'ReadWriteOncePod' ]]; then
          change_access_mode $1
        elif [[ "${1}" == 'quick' ]]; then 
          apply_quick_mode 
        elif [[ "${1}" == 'direct' ]]; then 
          apply_fio_direct
        elif [[ "${1}" == 'storage' ]]; then 
          change_storage_size "${GLOBAL_STORAGE_SIZE}"
        fi 
      fi 
      # change storageclass name from your-storageclass to $choice 
      set_storageclass "${choice}"
      # actual run kdbench and show up to stdout
      run_kdbench 
      # delete all objects for convenience purpose
      delete_pvc_after_job_done
  fi
}

choose_n_run_storageclasses_direct() {

# file mode (optional)
  if [[ "${1}" == 'file' ]]; then 
    run_kdbench $1 
    delete_pvc_after_job_done
    exit 0
  fi 
  # normal direct mode 
  local choice
  choice="$1"

    create_yaml_file
    # change storageclass name from your-storageclass to $1 
    set_storageclass "${choice}"
    # actual run kdbench and show up to stdout
    run_kdbench
    # delete all objects for convenience purpose
    delete_pvc_after_job_done
}

apply_quick_mode() {
  sed -i'' -e "s|# - name: KDBENCH_QUICK|- name: KDBENCH_QUICK|g" "$YAMLFILE"
  sed -i'' -e 's|#   value: "yes"|  value: "yes"|g' "$YAMLFILE"
}

apply_fio_direct() {
  sed -i'' -e "s|# - name: FIO_DIRECT|- name: FIO_DIRECT|g" "$YAMLFILE"
  sed -i'' -e 's|#   value: "1"|  value: "1"|g' "$YAMLFILE"
}

change_access_mode() {
  sed -i'' -e "s|ReadWriteOnce|$1|g" "$YAMLFILE"
}

change_storage_size() {
  sed -i'' -e "s|storage: 100|storage: $1|g" "$YAMLFILE"
}

create_yaml_file(){
  cp -r "$SCRACHFILE" "$YAMLFILE"
}

validate_kubectl() {
  if [[ -z "${KUBECTL:-}" ]]; then
    if hash kubectl 2>/dev/null; then
      KUBECTL=kubectl
    elif hash kubectl.exe  2>/dev/null; then
      KUBECTL=kubectl.exe
    else
      echo >&2 "kubectl is not installed"
      exit_err
    fi
  fi
}

run_storageclasses_interactive(){
  if [[ -t 1 &&  -z "${KDBENCH_IGNORE_FZF:-}" && "$(type fzf &>/dev/null; echo $?)" -eq 0 ]]; then
    if [[ "$#" -eq 0 ]]; then
      choose_n_run_storageclasses_interactive
    elif [[ "$#" -eq 1 ]]; then
      choose_n_run_storageclasses_interactive "${1}"
    else
      :
    fi 
  else
    get_storageclasses 
  fi 
}

main() {
  validate_kubectl
  
  # basic_run_interative
  if [[ "$#" -eq 0 ]]; then
    run_storageclasses_interactive
  # quick_mode_on
  elif [[ "${1}" == '-q' || "${1}" == '--quick' ]]; then
    run_storageclasses_interactive "quick" 
  # write_fio_direct_on
  elif [[ "${1}" == '-d' || "${1}" == '--direct' ]]; then
    run_storageclasses_interactive "direct"
  # reset to init stage   
  elif [[ "${1}" == 'reset' ]]; then
    rm -f $YAMLFILE $YAMLFILE-e
    $KUBECTL delete -f $SCRACHFILE 2>/dev/null || echo "Delete partly kdbench resource completely"
    $KUBECTL get po,pv,pvc -l fio=kdbench 
  elif [[ "${1}" == '-h' || "${1}" == '--help' ]]; then
    usage
  elif [[ "${1}" == '-a' ]]; then
    if [[ "${2}" == 'RWO' || "${2}" == 'rwo'  ]]; then
      :
    elif [[ "${2}" == 'ROX' || "${2}" == 'rox'  ]]; then
      run_storageclasses_interactive "ReadOnlyMany"
    elif [[ "${2}" == 'RWX' || "${2}" == 'rwx'  ]]; then
      run_storageclasses_interactive "ReadWriteMany"
    elif [[ "${2}" == 'RWOP' || "${2}" == 'rwop'  ]]; then
      run_storageclasses_interactive "ReadWriteOncePod"
    else 
      echo "error: check access_mode in RWO,ROX,RWX,RWOP"
    fi
  # storage size   
  elif [[ "${1}" == '-s' ]]; then
    if [[ "${2}" =~ ^-?[0-9]+$ ]]; then
      GLOBAL_STORAGE_SIZE="${2}"
      run_storageclasses_interactive "storage" 
    else
      echo "error: check storage class size number" >&2
      exit_err
    fi
  # From Yaml file 
  elif [[ "${1}" == '-f' ]]; then
    if [[ "$#" -eq 2 ]]; then
      GLOBAL_File_NAME="${2}"
      choose_n_run_storageclasses_direct "file" 
    else
      echo "error: check Yaml file name" >&2
      exit_err
    fi 
  # run_by_storage_name_manually 
  elif [[ "$#" -eq 1 ]]; then
    STORAGECLASSESLIST=$(kubectl get storageclasses -o=name | cut -c 29- 2>&1)
    if printf '%s\n' "${STORAGECLASSESLIST[@]}" | grep -q -F -x "$1" ; then
      choose_n_run_storageclasses_direct $1
    else
      echo "error: check storageclass name or flag option" >&2
      exit_err
    fi 
  elif [[ "$#" -gt 2 ]]; then
    echo "error: too many arguments" >&2
    usage
    exit_err
  else
    echo "error: unrecognized flag \"${1}\" & value \"${2}"\ >&2
    usage
    exit_err
  fi
}

main "$@"