#!/bin/bash
#
# Persistent Reservation (PR) test utility
#
# Author: qinwang@redhat.com
#

DEV=""
TESTCASE=0
KEY="0x123abc"
PAST_KEY="0x123abc"
INTERVAL=3
PR_TYPE="5"
OTHER_ARGS=""
PR_BIN=""
USAGE="
Usage:
  $0 -d <dev> -c <case_idx> [options]

Options:
  -d <dev>        Block device (e.g. /dev/sdX)
  -c <case_idx>   Test case index
  -k <key>        Reservation key (default: ${KEY})
  -p <pastkey>    Past/old key (for preempt)
  -t <pr_type>    Reservation type ()
  -o <args>       Other arguments (e.g. force, f)
  -i <interval>   Interval seconds (default: ${INTERVAL})
  -h              Show this help

Test cases:
  0: Read keys and reservation
  1: Read keys
  2: Read reservation
  3: Register        [-o force]
  4: Unregister
  5: Reserve
  6: Release
  7: Preempt         <new_key> <old_key>
  8: Clear
  9: Integration test
 10: Read capabilities and full status

Examples:
  # Register key 0x123abc
  $0 -d /dev/sdX -c 3 -k 0x123abc

  # Replace key with 0x456abc
  $0 -d /dev/sdX -c 3 -k 0x456abc -o f

  # Preempt with new key
  $0 -d /dev/sdX -c 7 -k 0x456abc -p 0x123abc

  # Integration test with interval
  $0 -d /dev/sdX -c 9 -i 1
"

DEFAULT_LOGDIR="$(dirname $0)/logs"
LOGDIR=${LOGDIR:-$DEFAULT_LOGDIR}
LOGFILE="$LOGDIR/$(basename $0 .sh)-$(date +%F-%H%M%S).log"


while getopts "hd:k:c:p:i:o:t:l:" opt; do
  case $opt in
  h)
    echo -e "$USAGE"
    exit 0
    ;;
  d)
    DEV="$OPTARG"
    ;;
  t)
    PR_TYPE="$OPTARG"
    ;;
  k)
    KEY="$OPTARG"
    ;;
  c)
    TESTCASE="$OPTARG"
    ;;
  p)
    PAST_KEY="$OPTARG"
    ;;
  i)
    INTERVAL="$OPTARG"
    ;;
  o)
    OTHER_ARGS="$OPTARG"
    ;;
  l)
    LOGFILE="$OPTARG"
    ;;
  ? | *)
    echo -e "$USAGE"
    echo "Unknown parameter"
    exit 1
    ;;
  esac
done

if [[ -z "${DEV}" || -z "$TESTCASE" ]]; then
  echo -e "Missing parameter or Wrong case.\n$USAGE\n"
  exit 1
fi

if [ ! -d "$(dirname ${LOGFILE})" ]; then
  mkdir -p "$(dirname $LOGFILE)"
fi

LOGLINK=$(dirname "$LOGFILE")/$(basename "$0" .sh)-latest.log
ln -srf "${LOGFILE}" "${LOGLINK}"
echo "Logging to ${LOGFILE} ${LOGLINK}"

exec &> >(tee -a >(awk '{ print strftime("%Y-%m-%d %H:%M:%S"), $0; fflush() }' >>"${LOGFILE}"))


# if [ "${DEV}" == "" ]; then
#   echo "No dev"
#   echo -e "${USAGE}"
#   exit 1
# fi

if [ "$PR_BIN" == "" ]; then
  if echo "${DEV}" | grep mapper; then
    PR_BIN="mpathpersist"
  else
    PR_BIN="sg_persist"
  fi
fi

echo "$PR_BIN : Dev:${DEV}  Testcase:${TESTCASE} Key:${KEY}"

pr_cmd() {
  echo "$@"
  eval "$@"
  ret=$?
  [ $ret -ne 0 ] && {
    echo "error:$@"
    exit $ret
  }
  return 0
}

pr_cmd_ignore_error() {
  echo "$@"
  eval "$@"
  return $?
  
}

pr_wait_for() {
  local cmd="$1"                  # command
  local wait_exit_code=${2:-0}    # exit code default 0（succeed）
  local timeout=${3:-60}          # timeout default 60
  local interval=${4:-1}          # interval default 5
  local fail_msg=${5:-"Timeout"}  # message
  local show_progress=${6:-true}  # display progess
  local wait_msg=${7:-""}         # wait message

  local elapsed=0
  local exit_code=0
  local display_progress=0

  if [ -n "$wait_msg" ]; then
    echo "Wait for $wait_msg"
  else
    echo "Wait for command: $cmd"
  fi

  while true; do
    eval "$cmd"; exit_code=$?
    
    if [ "$exit_code" -eq "$wait_exit_code" ]; then
      if [ "$display_progress" = "1" ]; then
          echo " "
      fi
      break
    fi

    if [ "$show_progress" = "true" ]; then
      # echo -n ".($exit_code)"
      echo -n "."
      display_progress=1
    fi

    sleep "$interval"
    elapsed=$((elapsed + interval))

    if [ "$elapsed" -ge "$timeout" ]; then
      if [ "$show_progress" = "true" ]; then
        echo " "
      fi
      echo "$fail_msg timeout after ${timeout}s"
      return 1
    fi
  done

  echo -e "Wait $wait_msg end($(date "+%H:%M:%S")) time:$elapsed\n"
  return 0
}
pr_read_key() {
  dev=${1:-${DEV}}
  pr_cmd "${PR_BIN} --in -k ${dev}"
}

pr_read_reservation() {
  dev=${1:-${DEV}}
  pr_cmd "${PR_BIN} --in -r ${dev}"
}

pr_read_full_status() {
  dev=${1:-${DEV}}
  pr_cmd "${PR_BIN} --in -s ${dev}"
}

pr_register_ignore() {
  # force register / replace current key
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  pr_cmd "${PR_BIN} --out --register-ignore --param-sark ${key} ${dev}"
}
pr_register() {
  # each node only one key
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  force=${3:-${OTHER_ARGS:-0}}

  if [ "${force}" -eq 1 ]; then
    pr_cmd "${PR_BIN} --out --register-ignore --param-sark ${key} ${dev}"
  else
    pr_cmd "${PR_BIN} --out --register --param-sark ${key} ${dev}"
  fi
  
}
pr_reserve() {
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  pr_cmd "${PR_BIN} --out --reserve --param-rk ${key} --prout-type ${PR_TYPE}  ${dev}"
}

pr_release() {
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  pr_cmd "${PR_BIN} --out --release --param-rk ${key} --prout-type ${PR_TYPE}  ${dev}"
}

pr_unregister() {
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  pr_cmd "${PR_BIN} --out --register --param-rk ${key} --param-sark=0 ${dev}"

}

pr_key_clear() {
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
    
  pr_cmd "${PR_BIN} --out --clear --param-rk=${key} ${dev}"
}

pr_preempt(){
  # It need to 2 hosts to test this.
  dev=${1:-${DEV}}
  new_key=${2:-${KEY}}
  old_key=${3:-${PAST_KEY}}
  pr_register_ignore "${dev}" "${new_key}"
  pr_cmd "${PR_BIN} --out --preempt --param-sark=${old_key} --param-rk=${new_key} --prout-type=${PR_TYPE} ${dev}"

}

pr_integration_test() {
  dev=${1:-${DEV}}
  key=${2:-${KEY}}
  interval=${3:-${INTERVAL}}
  pr_cmd pr_read_key "${dev}"
  pr_cmd pr_read_reservation "${dev}"

  pr_cmd pr_register "${dev}" "${key}"
  pr_cmd sleep "${interval}"
  pr_wait_for "pr_read_key |grep -q ${key}"

  pr_cmd pr_reserve "${dev}" "${key}"
  pr_cmd sleep "${interval}"
  pr_wait_for "pr_read_reservation |grep -q ${key}"
  
  pr_cmd pr_release "${dev}" "${key}"
  pr_cmd sleep "${interval}"
  pr_wait_for "pr_read_reservation |grep -q ${key}" 1  60 1

  pr_cmd pr_unregister "${dev}" "${key}"
  pr_cmd sleep "${interval}"
  pr_wait_for "pr_read_key | grep -q ${key}" 1 120 1
}

pr_read_capabilities(){
  dev=${1:-${DEV}}
  pr_cmd "${PR_BIN} --in --report-capabilities ${dev}"
}


case ${TESTCASE} in
0)
  echo "Read key and reservation..."
  pr_read_key "${DEV}"
  pr_read_reservation "${DEV}"
  ;;
1)
  echo "Read key..."
  pr_read_key "${DEV}"

  ;;
2)
  echo "Read reservation..."
  pr_read_reservation "${DEV}"
  ;;
3)
  echo "Register key..."
  pr_register "${DEV}" "${KEY}" "${OTHER_ARGS}"
  pr_read_key "${DEV}"
  ;;
4)
  echo "Unregister key ..."
  pr_unregister "${DEV}" "${KEY}"
  pr_read_key "${DEV}"
  ;;
5)
  echo "Reserve reservation ..."
  pr_reserve "${DEV}" "${KEY}"   
  pr_read_reservation "${DEV}"
  ;;
6)
  echo "Release reservation ..."
  pr_release  "${DEV}" "${KEY}"
  pr_read_reservation "${DEV}"
  ;;
7)
  echo "preempt ..."
  pr_preempt "${DEV}" "${KEY}"
  pr_read_key "${DEV}"         
  pr_read_reservation "${DEV}"
  ;;
8)
  echo "Clear ..."
  pr_key_clear "${DEV}" "${KEY}"
  pr_read_key   "${DEV}"
  pr_read_reservation "${DEV}"
  ;;
9)
  echo "Integration Test..."
  pr_integration_test "${DEV}" "${KEY}"
  ;;
10)
  echo "Read capabilities..."
  pr_read_capabilities "${DEV}"
  pr_read_full_status "${DEV}"
  ;;
?)
  echo "$USAGE"
  ;;
*)
  echo "Unknown testcase ${TESTCASE}"
  pr_cmd "${TESTCASE}"
  ;;
esac

# echo "0: "
# pr_cmd "${PR_BIN} --in -k ${dev}"
# pr_cmd "${PR_BIN} --in -r ${dev}"

# echo "1:register-key"
# pr_cmd "${PR_BIN} --out --register-ignore --param-sark 123aaa ${dev}"

# echo "2:read-key"
# pr_cmd "${PR_BIN}  --in -k ${dev}"

# echo "3:reserve"
# pr_cmd "${PR_BIN} --out --reserve --param-rk 123aaa --prout-type 5  ${dev}"

# echo "4:read-reservation"
# pr_cmd "${PR_BIN} --in -r ${dev}"

# echo "5:release"
# pr_cmd "${PR_BIN} --out --release --param-rk 123aaa --prout-type 5  ${dev}"

# echo "6:read-reservation"
# pr_cmd "${PR_BIN} --in -r ${dev}"

# echo "7:cancel-register"
# pr_cmd "${PR_BIN} --out --register-ignore --param-rk 123aaa --prout-type 5  ${dev}"

# echo "8:read-key"
# pr_cmd "${PR_BIN} --in -k ${dev}"
