#!/bin/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.
#
# Starts a Flume NG agent
#
# chkconfig: 345 90 10
# description: Flume NG agent
#
### BEGIN INIT INFO
# Provides:          flume-agent
# Required-Start:    $remote_fs
# Should-Start:
# Required-Stop:     $remote_fs
# Should-Stop:
# Default-Start:     3 4 5
# Default-Stop:      0 1 2 6
# Short-Description: Flume NG agent
### END INIT INFO

. /lib/lsb/init-functions

BIGTOP_DEFAULTS_DIR=${BIGTOP_DEFAULTS_DIR-/etc/default}
[ -n "${BIGTOP_DEFAULTS_DIR}" -a -r ${BIGTOP_DEFAULTS_DIR}/flume-agent ] && . ${BIGTOP_DEFAULTS_DIR}/flume-agent

# Autodetect JAVA_HOME if not defined
. /usr/lib/bigtop-utils/bigtop-detect-javahome

STATUS_RUNNING=0
STATUS_DEAD=1
STATUS_DEAD_AND_LOCK=2
STATUS_NOT_RUNNING=3

ERROR_PROGRAM_NOT_INSTALLED=5

FLUME_LOG_DIR=/var/log/flume
FLUME_CONF_DIR=/etc/flume/conf
FLUME_RUN_DIR=/var/run/flume
FLUME_HOME=/usr/lib/flume
FLUME_USER=flume

FLUME_LOCK_DIR="/var/lock/subsys/"
LOCKFILE="${FLUME_LOCK_DIR}/flume-agent"
desc="Flume NG agent daemon"

FLUME_CONF_FILE=${FLUME_CONF_FILE:-${FLUME_CONF_DIR}/flume.conf}
EXEC_PATH=/usr/bin/flume-ng
FLUME_PID_FILE=${FLUME_RUN_DIR}/flume-agent.pid

# These directories may be tmpfs and may or may not exist
# depending on the OS (ex: /var/lock/subsys does not exist on debian/ubuntu)
for dir in "$FLUME_RUN_DIR" "$FLUME_LOCK_DIR"; do
  [ -d "${dir}" ] || install -d -m 0755 -o $FLUME_USER -g $FLUME_USER ${dir}
done

DEFAULT_FLUME_AGENT_NAME="agent"
FLUME_AGENT_NAME=${FLUME_AGENT_NAME:-${DEFAULT_FLUME_AGENT_NAME}}
FLUME_SHUTDOWN_TIMEOUT=${FLUME_SHUTDOWN_TIMEOUT:-60}

# Update variables if service request is followed by a service name
# e.g. service flume-agent start|stop|status|restart service-name
# Starts a flume agent with the name flume-service-name
# Expects flume-service-name.conf file in FLUME_CONF_DIR
setvariables() {
  AGENT_NAME=$1
  LOCKFILE="${FLUME_LOCK_DIR}/flume-${AGENT_NAME}"
  FLUME_CONF_FILE=${FLUME_CONF_DIR}/flume-${AGENT_NAME}.conf
  FLUME_PID_FILE=${FLUME_RUN_DIR}/flume-${AGENT_NAME}.pid
  DEFAULT_FLUME_AGENT_NAME=${AGENT_NAME}
  FLUME_AGENT_NAME=${DEFAULT_FLUME_AGENT_NAME}
}

start() {
  [ -x $exec ] || exit $ERROR_PROGRAM_NOT_INSTALLED

  checkstatus
  status=$?
  if [ "$status" -eq "$STATUS_RUNNING" ]; then
    return 0
  fi

  log_success_msg "Starting $desc $FLUME_AGENT_NAME: "
  /bin/su -s /bin/bash -c "/bin/bash -c 'echo \$\$ >${FLUME_PID_FILE} && exec ${EXEC_PATH} agent --conf $FLUME_CONF_DIR --conf-file $FLUME_CONF_FILE --name $FLUME_AGENT_NAME >>${FLUME_LOG_DIR}/flume-${FLUME_AGENT_NAME}.out 2>&1' &" $FLUME_USER
  RETVAL=$?
  [ $RETVAL -eq 0 ] && touch $LOCKFILE
  return $RETVAL
}

stop() {
  if [ ! -e $FLUME_PID_FILE ]; then
    log_failure_msg "Flume agent $FLUME_AGENT_NAME is not running"
    return 0
  fi

  log_success_msg "Stopping $desc $FLUME_AGENT_NAME: "

  FLUME_PID=`cat $FLUME_PID_FILE`
  if [ -n $FLUME_PID ]; then
    kill -TERM ${FLUME_PID} &>/dev/null
    for i in `seq 1 ${FLUME_SHUTDOWN_TIMEOUT}` ; do
      kill -0 ${FLUME_PID} &>/dev/null || break
      sleep 1
    done
    kill -KILL ${FLUME_PID} &>/dev/null
  fi
  rm -f $LOCKFILE $FLUME_PID_FILE
  return 0
}

restart() {
  stop
  start
}

checkstatus(){
  pidofproc -p $FLUME_PID_FILE java > /dev/null
  status=$?

  case "$status" in
    $STATUS_RUNNING)
      log_success_msg "Flume NG agent $FLUME_AGENT_NAME is running"
      ;;
    $STATUS_DEAD)
      log_failure_msg "Flume NG agent $FLUME_AGENT_NAME is dead and pid file exists"
      ;;
    $STATUS_DEAD_AND_LOCK)
      log_failure_msg "Flume NG agent $FLUME_AGENT_NAME is dead and lock file exists"
      ;;
    $STATUS_NOT_RUNNING)
      log_failure_msg "Flume NG agent $FLUME_AGENT_NAME is not running"
      ;;
    *)
      log_failure_msg "Flume NG agent $FLUME_AGENT_NAME status is unknown"
      ;;
  esac
  return $status
}

condrestart(){
  [ -e ${LOCKFILE} ] && restart || :
}

#
# If user doesn't provide a flume agent name to start,
# will attempt to start all agents based on conf files in ${FLUME_CONF_DIR} 
#
startall() {
  run_functions_on_conf start
}

#
# If user doesn't provide a flume agent name to stop,
# will attempt to stop all agents based on conf files in ${FLUME_CONF_DIR} 
#
stopall() {
  run_functions_on_conf stop
}

#
# If user doesn't provide a flume agent name to restart,
# will attempt to restart all agents based on conf files in ${FLUME_CONF_DIR} 
#
restartall() {
  run_functions_on_conf restart
}

#
# If user doesn't provide a flume agent name to perform cond restart,
# will attempt to do condrestart all agents based on conf files in ${FLUME_CONF_DIR} 
#
condrestartall() {
  run_functions_on_conf condrestart
}

#
# If user doesn't provide a flume agent name to check the status,
# status of all agents based on conf files in ${FLUME_CONF_DIR} is checked
#
checkallstatus() {
  run_functions_on_conf checkstatus
}

#
# Common function to perform user action on all flume conf files
#
run_functions_on_conf() {
  if [ -f ${FLUME_CONF_DIR}/flume.conf ]; then
    echo $FLUME_AGENT_NAME
    $1
  fi
  agent_conf_pattern="${FLUME_CONF_DIR}/flume-.*\.conf$"
  for f in ${FLUME_CONF_DIR}/*
  do
    if [ $(echo $f | grep -e ${agent_conf_pattern}) ]; then
      conf_file=${f%.*}
      file_name=${conf_file##*/}
      agent_name=${file_name#*-}
      echo $agent_name
      setvariables $agent_name
      $1
    fi
  done
}

#
# Logic to update variables if service start|stop|... is made with a flume agent name
#
if [ "$#" -eq 2 ]; then
  setvariables $2
fi

case "$1" in
  start)
    if [ "$#" -eq 2 ];then
      start
    else
      startall
    fi
    ;;
  stop)
    if [ "$#" -eq 2 ]; then
      stop
    else
      stopall
    fi
    ;;
  status)
    if [ "$#" -eq 2 ]; then
      checkstatus
    else
      checkallstatus
    fi
    ;;
  restart)
    if [ "$#" -eq 2 ]; then
      restart
    else
      restartall
    fi
    ;;
  condrestart|try-restart)
    if [ "$#" -eq 2 ]; then
      condrestart
    else
      condrestartall
    fi
    ;;
  *)
    echo $"Usage: $0 {start|stop|status|restart|try-restart|condrestart}"
    RETVAL=1
esac

exit $RETVAL
