kyuubi/bin/kyuubi

#!/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.
#

## Kyuubi Server Main Entrance
CLASS="org.apache.kyuubi.server.KyuubiServer"

function usage() {
  echo "Usage: bin/kyuubi command"
  echo "  commands:"
  echo "    start        - Run a Kyuubi server as a daemon"
  echo "    restart      - Restart Kyuubi server as a daemon"
  echo "    run          - Run a Kyuubi server in the foreground"
  echo "    stop         - Stop the Kyuubi daemon"
  echo "    kill         - Stop the Kyuubi daemon forcefully"
  echo "    status       - Show status of the Kyuubi daemon"
  echo "    -h | --help  - Show this help message"
}

if [[ "$*" = *--help ]] || [[ "$*" = *-h ]]; then
  usage
  exit 0
fi

# if for some reason the shell doesn't have $USER defined
# (e.g., ssh'd in to execute a command)
# let's get the effective username and use that
USER=${USER:-$(id -nu)}

function kyuubi_logo() {
  source ${KYUUBI_HOME}/bin/kyuubi-logo
}

function kyuubi_rotate_log() {
  log=$1;

  if [[ -z ${KYUUBI_MAX_LOG_FILES} ]]; then
    num=5
  elif [[ ${KYUUBI_MAX_LOG_FILES} -gt 0 ]]; then
    num=${KYUUBI_MAX_LOG_FILES}
  else
    echo "Error: KYUUBI_MAX_LOG_FILES must be a positive number, but got ${KYUUBI_MAX_LOG_FILES}"
    # shellcheck disable=SC2242
    exit -1
  fi

  if [ -f "$log" ]; then # rotate logs
  while [ ${num} -gt 1 ]; do
    prev=`expr ${num} - 1`
    [ -f "$log.$prev" ] && mv "$log.$prev" "$log.$num"
    num=${prev}
  done
  mv "$log" "$log.$num";
  fi
}

if [ -z "${KYUUBI_HOME}" ]; then
  KYUUBI_HOME="$(cd "`dirname "$0"`"/..; pwd)"
fi

if [[ $1 == "start" ]] || [[ $1 == "run" ]]; then
  . "${KYUUBI_HOME}/bin/load-kyuubi-env.sh"
else
  . "${KYUUBI_HOME}/bin/load-kyuubi-env.sh" -s
fi

if [[ -z ${JAVA_HOME} ]]; then
  echo "Error: JAVA_HOME IS NOT SET! CANNOT PROCEED."
  exit 1
fi

RUNNER="${JAVA_HOME}/bin/java"

## Find the Kyuubi Jar
if [[ -z "$KYUUBI_JAR_DIR" ]]; then
  KYUUBI_JAR_DIR="$KYUUBI_HOME/jars"
  if [[ ! -d ${KYUUBI_JAR_DIR} ]]; then
  echo -e "\nCandidate Kyuubi lib $KYUUBI_JAR_DIR doesn't exist, searching development environment..."
    KYUUBI_JAR_DIR="$KYUUBI_HOME/kyuubi-assembly/target/scala-${KYUUBI_SCALA_VERSION}/jars"
  fi
fi

## Find the WebUI dist
if [[ -z "$KYUUBI_WEBUI_DIR" ]]; then
  KYUUBI_WEBUI_DIR="$KYUUBI_HOME/web-ui"
  if [[ ! -d ${KYUUBI_WEBUI_DIR} ]]; then
  echo -e "\nCandidate Kyuubi WebUI dist $KYUUBI_WEBUI_DIR doesn't exist, searching development environment..."
    KYUUBI_WEBUI_DIR="$KYUUBI_HOME/kyuubi-server/web-ui"
  fi
fi

KYUUBI_CLASSPATH="${KYUUBI_CONF_DIR}:${KYUUBI_WEBUI_DIR}:${KYUUBI_JAR_DIR}/*:${HADOOP_CONF_DIR}"
if [[ -n ${YARN_CONF_DIR} ]]; then
  KYUUBI_CLASSPATH="${KYUUBI_CLASSPATH}:${YARN_CONF_DIR}"
fi

if [[ "$1" =~ ^(start|restart|run|stop|kill|status)$ ]]; then
    command=$1
    shift
fi

cmd="${RUNNER} ${KYUUBI_JAVA_OPTS} -cp ${KYUUBI_CLASSPATH} $CLASS $@"

pid="${KYUUBI_PID_DIR}/kyuubi-$USER-$CLASS.pid"

function start_kyuubi() {
  if [[ ! -w ${KYUUBI_PID_DIR} ]]; then
    echo "${USER} does not have 'w' permission to ${KYUUBI_PID_DIR}"
    exit 1
  fi

  if [[ ! -w ${KYUUBI_LOG_DIR} ]]; then
    echo "${USER} does not have 'w' permission to ${KYUUBI_LOG_DIR}"
    exit 1
  fi

  if [ -f "$pid" ]; then
    TARGET_ID="$(cat "$pid")"
    if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then
      echo "$CLASS running as process $TARGET_ID  Stop it first."
      exit 1
    fi
  fi

  log="${KYUUBI_LOG_DIR}/kyuubi-$USER-$CLASS-$HOSTNAME.out"
  kyuubi_rotate_log ${log}

  echo "Starting $CLASS, logging to $log"
  nohup nice -n "${KYUUBI_NICENESS:-0}" ${cmd} >> ${log} 2>&1 < /dev/null &
  newpid="$!"

  echo "$newpid" > "$pid"

  # Poll for up to 5 seconds for the java process to start
  for i in {1..10}
  do
    if [[ $(ps -p "$newpid" -o comm=) =~ "java" ]]; then
       break
    fi
    sleep 0.5
  done

  sleep 2
  # Check if the process has died; in that case we'll tail the log so the user can see
  if [[ ! $(ps -p "$newpid" -o comm=) =~ "java" ]]; then
    echo "Failed to launch: ${cmd}"
    tail -2 "$log" | sed 's/^/  /'
    echo "Full log in $log"
  else
    echo "Welcome to"
    kyuubi_logo
  fi
}

function run_kyuubi() {
  echo "Starting $CLASS"
  exec nice -n "${KYUUBI_NICENESS:-0}" ${cmd}
}

function stop_kyuubi() {
  if [ -f ${pid} ]; then
    TARGET_ID="$(cat "$pid")"
    if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then
      echo "Stopping $CLASS"
      kill "$TARGET_ID"
      for i in {1..20}
      do
        sleep 0.5
        if [[ ! $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then
          rm -f "$pid"
          break
        fi
      done

      if [ -f ${pid} ]; then
        echo "Failed to stop kyuubi after 10 seconds, try 'kill -9 ${TARGET_ID}' forcefully"
        kill -9 ${TARGET_ID}
        for i in {1..20}
        do
          sleep 0.5
          if [[ ! $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then
            rm -f "$pid"
            break
          fi
        done
        if [ -f ${pid} ]; then
          echo "Failed to stop kyuubi forcefully after 10 seconds"
          exit 1
        else
          kyuubi_logo
          echo "Bye!"
        fi
      else
        kyuubi_logo
        echo "Bye!"
      fi
    else
      echo "no $CLASS to stop"
    fi
  else
    echo "no $CLASS to stop"
  fi
}

function kill_kyuubi() {
  if [ -f ${pid} ]; then
    TARGET_ID="$(cat "$pid")"
    echo "Stopping $CLASS forcefully."
    kill -9 $TARGET_ID
    kyuubi_logo
    echo "Bye!"
  else
    echo "no $CLASS to stop"
  fi
}

function check_kyuubi() {
  if [[ -f ${pid} ]]; then
    TARGET_ID="$(cat "$pid")"
    if [[ $(ps -p "$TARGET_ID" -o comm=) =~ "java" ]]; then
      echo "Kyuubi is running (pid: $TARGET_ID)"
    else
      echo "Kyuubi is not running"
    fi
  else
    echo "Kyuubi is not running"
  fi
}

case $command in
  (start | "")
    start_kyuubi
    ;;

  (restart)
    echo "Restarting Kyuubi"
    stop_kyuubi
    start_kyuubi
    ;;

  (run)
    run_kyuubi
    ;;

  (stop)
    stop_kyuubi
    ;;

  (kill)
    kill_kyuubi
    ;;

  (status)
    check_kyuubi
    ;;

  (*)
    usage
    ;;
esac