#!/usr/bin/env bash
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#

######################################
# Proton common functions
######################################

function proton_error() {
  echo "$*" 1>&2
}

function proton_status_daemon() {
   # 0 = program is running
   # 1 = dead, but still a pid (2)
   # 2 = not running

  local pid_file=$1
  local daemon_name=$2
  shift

  local pid
  local pspid

  if [[ -f "${pid_file}" ]]; then
    pid=$(cat "${pid_file}")
    if pspid=$(ps -o args= -p"${pid}" 2>/dev/null); then
      # this is to check that the running process we found is actually the same daemon that we're interested in
      if [[ ${pspid} =~ -Dproc_${daemon_name} ]]; then
        return 0
      fi
    fi

    return 1
  fi
  return 2
}

function proton_start_daemon() {
  local command=$1
  local class=$2
  local pidfile=$3
  shift 3

  # this is for the non-daemon pid creation
  #shellcheck disable=SC2086
  echo $$ > "${pidfile}" 2>/dev/null
  if [[ $? -gt 0 ]]; then
    proton_error "Error: Cannot write ${command} pid ${pidfile}."
  fi

  export CLASSPATH
  exec "${JAVA}" "-Dproc_${command}" ${PROTON_OPTS} "${class}" "$@"
}

function proton_start_daemon_wrapper() {
  local daemonname=$1
  local class=$2
  local pidfile=$3
  local outfile=$4
  shift 4

  local counter

  proton_start_daemon "${daemonname}" \
    "$class" \
    "${pidfile}" \
    "$@" >> "${outfile}" 2>&1 < /dev/null &

  # we need to avoid a race condition here
  # so let's wait for the fork to finish
  # before overriding with the daemonized pid
  (( counter=0 ))
  while [[ ! -f ${pidfile} && ${counter} -le 5 ]]; do
    sleep 1
    (( counter++ ))
  done

  # This is for daemon pid creation, $! is the process ID of the last job that was background.
  # shellcheck disable=SC2086
  echo $! > "${pidfile}" 2>/dev/null
  if [[ $? -gt 0 ]]; then
    proton_error "Error:  Cannot write ${daemonname} pid ${pidfile}."
  fi

  # shellcheck disable=SC2086
  disown %+ >/dev/null 2>&1
  if [[ $? -gt 0 ]]; then
    proton_error "Error: Cannot disconnect ${daemonname} process $!"
  fi
  sleep 1

  # capture the ulimit output
  ulimit -a >> "${outfile}" 2>&1

  # shellcheck disable=SC2086
  if ! ps -p $! >/dev/null 2>&1; then
    return 1
  fi
  return 0
}

function wait_process_to_die_or_timeout() {
  local pid=$1
  local timeout=$2

  # Normalize timeout
  # Round up or down
  timeout=$(printf "%.0f\n" "${timeout}")
  if [[ ${timeout} -lt 1  ]]; then
    # minimum 1 second
    timeout=1
  fi

  # Wait to see if it's still alive
  for (( i=0; i < "${timeout}"; i++ ))
  do
    if kill -0 "${pid}" > /dev/null 2>&1; then
      sleep 1
    else
      break
    fi
  done
}

function proton_stop_daemon() {
  local cmd=$1
  local pidfile=$2
  shift 2

  local pid
  local cur_pid

  if [[ -f "${pidfile}" ]]; then
    pid=$(cat "$pidfile")

    kill "${pid}" >/dev/null 2>&1

    wait_process_to_die_or_timeout "${pid}" "${PROTON_STOP_TIMEOUT}"

    if kill -0 "${pid}" > /dev/null 2>&1; then
      proton_error "WARNING: ${cmd} did not stop gracefully after ${PROTON_STOP_TIMEOUT} seconds: Trying to kill with kill -9"
      kill -9 "${pid}" >/dev/null 2>&1
    fi
    wait_process_to_die_or_timeout "${pid}" "${PROTON_STOP_TIMEOUT}"
    if ps -p "${pid}" > /dev/null 2>&1; then
      proton_error "ERROR: Unable to kill ${pid}"
    else
      cur_pid=$(cat "$pidfile")
      if [[ "${pid}" = "${cur_pid}" ]]; then
        rm -f "${pidfile}" >/dev/null 2>&1
      else
        proton_error "WARNING: pid has changed for ${cmd}, skip deleting pid file"
      fi
    fi
  fi
}

######################################
# Proton bash main entrypoint.
######################################

# Environment Variables:
#
#   JAVA_HOME                    The java implementation to use.  Overrides JAVA_HOME.
#
#   PROTON_CLASSPATH             Extra Java CLASSPATH entries.
#
#   PROTON_OPTS                  Extra Java runtime options.
#
#   PROTON_METASERVER_OPTS       Metaserver extra Java runtime options.
#
#   PROTON_DATASERVER_OPTS       Dataserver extra Java runtime options.
#
#   PROTON_CONF_DIR              Alternate conf dir. Default is ${PROTON_HOME}/conf.
#
#   PROTON_LOG_DIR               Alternate log dir. Default is ${PROTON_HOME}/logs.
#
#   PROTON_PID_DIR               Alternate pid dir. Default is ${PROTON_HOME}/logs.
#

bin=`dirname "$0"`
bin=`cd "$bin">/dev/null; pwd`

show_usage() {
  echo "Usage: proton [<options>] <command> [<args>]"
  echo "Options:"
  echo "  --help or -h                        Print this help message"
  echo "  --conf                              Proton config directory which includes proton.ini and log4j.properties"
  echo "  --daemon start|stop|status          Start, stop or status the daemon server process (such as meta server or data server)"
  echo "  --loglevel DEBUG|INFO|WARN|ERROR    Set the log level"
  echo ""
  echo "Commands:"
  echo "Some commands take arguments. Pass no args or -h for usage."
  echo "  cbench                              Run the cache benchmark."
  echo "  dataserver                          Run an Proton DataServer node."
  echo "  evict                               Evict all cached blocks of the specified file or all sub files if the path is a directory."
  echo "  fsck                                Verify the consistence between object storage and proton servers."
  echo "  fio                                 Run filesystem IO benchmark."
  echo "  load                                Load meta data or blocks from object storage and HDFS into proton clusters."
  echo "  metaserver                          Run an Proton MetaServer node."
  echo "  nnbench                             Run nnbench test."
  echo "  report                              Report Proton information and statistics."
  echo "  showblk                             Print the block information for given file."
  echo "  sync                                Sync the path between object storage and proton servers."
  echo "  version                             Print the version."
}

# if no args specified, show usage
if [ $# = 0 ]; then
  show_usage
  exit 1
fi

# Now having JAVA_HOME defined is required
if [ -z "$JAVA_HOME" ]; then
    echo "Error: JAVA_HOME is not set, Proton requires Java 1.8 or later."
    exit 1
fi
JAVA=$JAVA_HOME/bin/java

# The root of the proton installation
if [[ -z "$PROTON_HOME" ]]; then
  export PROTON_HOME=`dirname "$bin"`
fi

# Parse the options.
PROTON_CONF_DIR="${PROTON_CONF_DIR:-$PROTON_HOME/conf}"
PROTON_DAEMON_MODE=""
while true; do
  case $1 in
    --help|-help|-h|help|--h)
      show_usage
      exit 0
    ;;

    --conf)
      shift
      PROTON_CONF_DIR=$1
      shift
      if [[ ! -d "${PROTON_CONF_DIR}" ]]; then
        echo "Error: Cannot find the proton config directory \"${PROTON_CONF_DIR}\""
        exit 1
      fi
    ;;

    --daemon)
      shift
      PROTON_DAEMON_MODE=$1
      shift
      if [[ -z "${PROTON_DAEMON_MODE}" || ! "${PROTON_DAEMON_MODE}" =~ ^st(art|op|atus)$ ]]; then
        echo "Error: Invalid \"${PROTON_DAEMON_MODE}\", --daemon must be followed by either 'start', 'stop', or 'status'."
        exit 1
      fi
    ;;

    --loglevel)
      shift
      PROTON_LOG_LEVEL=$1
      shift
      if [[ -z "${PROTON_LOG_LEVEL}" || ! "${PROTON_LOG_LEVEL}" =~ ^(DEBUG|INFO|WARN|ERROR)$ ]]; then
        echo "Error: Invalid \"${PROTON_LOG_LEVEL}\", --log-level must be followed by either 'DEBUG', 'INFO', 'WARN', or 'ERROR'"
        exit 1
      fi
    ;;

    *)
      break
    ;;

  esac
done

# Get arguments
COMMAND=$1
shift

# Figure out which class to run
SUPPORT_DAEMON=false
case $COMMAND in
  cbench)
    CLASS="io.proton.cli.cache.CacheBench"
  ;;
  dataserver)
    CLASS="io.proton.core.main.DataServer"
    SUPPORT_DAEMON=true
    PROTON_OPTS="${PROTON_OPTS} -Dproton.log.file=proton-${COMMAND}.log"
    if [[ "${PROTON_DATASERVER_OPTS}" != "" ]]; then
          PROTON_OPTS="${PROTON_OPTS} ${PROTON_DATASERVER_OPTS}"
    fi
  ;;
  evict)
    CLASS="io.proton.cli.cmd.Evict"
  ;;
  fio)
    CLASS="io.proton.cli.fio.FIO"
  ;;
  fsck)
    CLASS="io.proton.cli.cmd.Fsck"
  ;;
  load)
    CLASS="io.proton.cli.cmd.Load"
  ;;
  metaserver)
    CLASS="io.proton.core.main.MetaServer"
    SUPPORT_DAEMON=true
    PROTON_OPTS="${PROTON_OPTS} -Dproton.log.file=proton-${COMMAND}.log"
    if [[ "${PROTON_METASERVER_OPTS}" != "" ]]; then
          PROTON_OPTS="${PROTON_OPTS} ${PROTON_METASERVER_OPTS}"
    fi
  ;;
  nnbench)
    CLASS="io.proton.cli.nnbench.NNBench"
  ;;
  report)
    CLASS="io.proton.cli.cmd.ReportInfo"
  ;;
  showblk)
    CLASS="io.proton.cli.cmd.ShowBlockInfo"
  ;;
  sync)
    CLASS="io.proton.cli.cmd.Sync"
  ;;
  version)
    CLASS="io.proton.common.VersionInfo"
  ;;
  *)
    proton_error "Error: Cannot find the command \"${COMMAND}\""
    exit 1
  ;;
esac

if [[ "${PROTON_DAEMON_MODE}" != "" && "${SUPPORT_DAEMON}" = "false" ]]; then
  proton_error "Error: Command \"${COMMAND}\" does not support --daemon option."
  exit 1
fi

# Check the existence of hadoop bundle jar.
HADOOP_MAJOR_VERSION=${HADOOP_MAJOR_VERSION:-3}
HADOOP_BUNDLE_JAR=$PROTON_HOME/plugins/hadoop${HADOOP_MAJOR_VERSION}/*.jar
HADOOP_BUNDLE_JAR=`ls ${HADOOP_BUNDLE_JAR}`
if [ ! -f $HADOOP_BUNDLE_JAR ]; then
  echo "Hadoop bundle jar does not exist: ${HADOOP_BUNDLE_JAR}"
  exit 1
fi

# Check the existence of web ui pages.
PROTON_WEBUI_PATH=${PROTON_WEBUI_PATH}
if [[ -z "$PROTON_WEBUI_PATH" ]]; then
  PROTON_WEBUI_PATH=${PROTON_HOME}
fi

# Now having HADOOP_HOME defined is required
if [ -z "$HADOOP_HOME" ]; then
    echo "Error: HADOOP_HOME is not set, Proton requires hadoop2.x or later"
    exit 1
fi
export LD_LIBRARY_PATH="${HADOOP_HOME}/lib/native"

# CLASSPATH initially contains $PROTON_CONF_DIR
CLASSPATH="${PROTON_CONF_DIR}:$HADOOP_BUNDLE_JAR:${PROTON_WEBUI_PATH}"
if [ "$PROTON_CLASSPATH" != "" ]; then
  CLASSPATH=${CLASSPATH}:$PROTON_CLASSPATH
fi
CLASSPATH=${CLASSPATH}:$JAVA_HOME/lib/tools.jar
for f in $PROTON_HOME/lib/*.jar; do
    CLASSPATH=${CLASSPATH}:$f;
done
CLASSPATH="${CLASSPATH}:`${HADOOP_HOME}/bin/hadoop classpath`"
export CLASSPATH

# Ensure the existence of $PROTON_LOG_DIR directory.
PROTON_LOG_DIR=${PROTON_LOG_DIR:-$PROTON_HOME/logs}
mkdir -p $PROTON_LOG_DIR

# Ensure the existence of $PROTON_PID_DIR directory.
PROTON_PID_DIR=${PROTON_PID_DIR:-$PROTON_LOG_DIR}
mkdir -p $PROTON_PID_DIR

# Configure $PROTON_OPTS
if [[ "$PROTON_LOG_LEVEL" != "" ]]; then
  PROTON_OPTS="${PROTON_OPTS} -Dproton.log.level=${PROTON_LOG_LEVEL}"
fi

# Execute the command, according to the daemon mode.
daemon_pidfile=${PROTON_PID_DIR}/proton-${COMMAND}.pid
daemon_outfile=${PROTON_LOG_DIR}/proton-${COMMAND}.out
case ${PROTON_DAEMON_MODE} in
  status)
    proton_status_daemon "${daemon_pidfile}"
    if [[ $? == 0 ]]; then
      echo "${COMMAND} is running as process $(cat "${daemon_pidfile}")."
      exit 0
    else
      echo "${COMMAND} is not running."
      exit 1
    fi
  ;;

  start)
    proton_status_daemon "${daemon_pidfile}"
    if [[ $? == 0  ]]; then
      proton_error "${COMMAND} is running as process $(cat "${daemon_pidfile}").  Stop it first and ensure ${daemon_pidfile} file is empty before retry."
      exit 1
    else
      # stale pid file, so just remove it and continue on
      rm -f "${daemon_pidfile}" >/dev/null 2>&1
    fi

    proton_start_daemon_wrapper "${COMMAND}" "${CLASS}" "${daemon_pidfile}" "${daemon_outfile}" "$@"
  ;;

  stop)
    proton_stop_daemon "${COMMAND}" "${daemon_pidfile}"
    exit $?
  ;;

  *)
    # Redirect the log into console output.
    PROTON_OPTS="${PROTON_OPTS} -Dproton.root.logger=${PROTON_LOG_LEVEL:-INFO},console"
    PROTON_PROPS="${PROTON_PROPS}"
    exec "$JAVA" -Dproc_$COMMAND -XX:OnOutOfMemoryError="kill -9 %p" $PROTON_OPTS $CLASS $PROTON_PROPS "$@"
  ;;
esac
