#!/bin/sh
### BEGIN INIT INFO
# Provides:          td-agent
# Required-Start:    $network $local_fs
# Required-Stop:
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: data collector for Treasure Data
# Description:       Treasure Data Support <support@treasure-data.com>
### END INIT INFO

set -e

export PATH=/sbin:/usr/sbin:/bin:/usr/bin

TD_AGENT_NAME=td-agent
TD_AGENT_HOME=/opt/td-agent
TD_AGENT_DEFAULT=/etc/default/td-agent
TD_AGENT_USER=td-agent
TD_AGENT_GROUP=td-agent
TD_AGENT_RUBY=/opt/td-agent/embedded/bin/ruby
TD_AGENT_BIN_FILE=/usr/sbin/td-agent
TD_AGENT_LOG_FILE=/var/log/td-agent/td-agent.log
TD_AGENT_PID_FILE=/var/run/td-agent/td-agent.pid
TD_AGENT_OPTIONS="--use-v1-config"

# timeout can be overridden from /etc/default/td-agent
STOPTIMEOUT=120

# Read configuration variable file if it is present
if [ -f "${TD_AGENT_DEFAULT}" ]; then
  . "${TD_AGENT_DEFAULT}"
fi

if [ -n "${NAME}" ]; then
  # backward compatibility with omnibus-td-agent <= 2.2.0. will be deleted from future release.
  echo "Warning: Declaring \$NAME in ${TD_AGENT_DEFAULT} for customizing \$PIDFILE has been deprecated. Use \$TD_AGENT_PID_FILE instead." 1>&2
  TD_AGENT_PID_FILE="//var/run/${NAME}/${NAME}.pid"
fi

if [ -n "${DAEMON_ARGS}" ]; then
# TODO: Show warning on use of `DAEMON_ARGS`
# echo "Warning: Declaring \$DAEMON_ARGS in ${TD_AGENT_DEFAULT} has been deprecated. Use \$TD_AGENT_OPTIONS instead." 1>&2
  TD_AGENT_OPTIONS="${DAEMON_ARGS} ${TD_AGENT_OPTIONS}"
fi

# Arguments to run the daemon with
TD_AGENT_ARGS="${TD_AGENT_ARGS:-${TD_AGENT_BIN_FILE} --log ${TD_AGENT_LOG_FILE} ${TD_AGENT_OPTIONS}}"
START_STOP_DAEMON_ARGS=""

# Exit if the package is not installed
[ -x "${TD_AGENT_RUBY}" ] || exit 0

# Source function library.
# Define LSB log_* functions.
# Depend on lsb-base (>= 3.0-6) to ensure that this file is present.
. /lib/lsb/init-functions

# Check the user
if [ -n "${TD_AGENT_USER}" ]; then
  if ! getent passwd | grep -q "^${TD_AGENT_USER}:"; then
    echo "$0: user for running ${TD_AGENT_NAME} doesn't exist: ${TD_AGENT_USER}" >&2
    exit 1
  fi
  mkdir -p "$(dirname "${TD_AGENT_PID_FILE}")"
  chown -R "${TD_AGENT_USER}" "$(dirname "${TD_AGENT_PID_FILE}")"
  START_STOP_DAEMON_ARGS="${START_STOP_DAEMON_ARGS} -c ${TD_AGENT_USER}"
fi

if [ -n "${TD_AGENT_GROUP}" ]; then
  if ! getent group | grep -q "^${TD_AGENT_GROUP}:"; then
    echo "$0: group for running ${TD_AGENT_NAME} doesn't exist: ${TD_AGENT_GROUP}" >&2
    exit 1
  fi
  START_STOP_DAEMON_ARGS="${START_STOP_DAEMON_ARGS} --group ${TD_AGENT_GROUP}"
fi

if [ -n "${TD_AGENT_PID_FILE}" ]; then
  mkdir -p "$(dirname "${TD_AGENT_PID_FILE}")"
  chown -R "${TD_AGENT_USER}" "$(dirname "${TD_AGENT_PID_FILE}")"
  TD_AGENT_ARGS="${TD_AGENT_ARGS} --daemon ${TD_AGENT_PID_FILE}"
fi

# 2012/04/17 Kazuki Ohta <k@treasure-data.com>
# Use jemalloc to avoid memory fragmentation
if [ -f "${TD_AGENT_HOME}/embedded/lib/libjemalloc.so" ]; then
  export LD_PRELOAD="${TD_AGENT_HOME}/embedded/lib/libjemalloc.so"
fi

#
# Function that starts the daemon/service
#
do_start() {
  # Set Max number of file descriptors for the safety sake
  # see http://docs.fluentd.org/en/articles/before-install
  ulimit -n 65536 1>/dev/null 2>&1 || true

  # Return
  #   0 if daemon has been started
  #   1 if daemon was already running
  #   2 if daemon could not be started
  start-stop-daemon --start --quiet --pidfile "${TD_AGENT_PID_FILE}" --exec "${TD_AGENT_RUBY}" \
    ${START_STOP_DAEMON_ARGS} --test > /dev/null \
    || return 1
  start-stop-daemon --start --quiet --pidfile "${TD_AGENT_PID_FILE}" --exec "${TD_AGENT_RUBY}" \
    ${START_STOP_DAEMON_ARGS} -- ${TD_AGENT_ARGS} \
    || return 2
  # Add code here, if necessary, that waits for the process to be ready
  # to handle requests from services started subsequently which depend
  # on this one.  As a last resort, sleep for some time.
}

#
# Function that stops the daemon/service
#
do_stop() {
  # Return
  #   0 if daemon has been stopped
  #   1 if daemon was already stopped
  #   2 if daemon could not be stopped
  #   other if a failure occurred
  local RETVAL=0
  start-stop-daemon --stop --quiet --retry="TERM/${STOPTIMEOUT}/KILL/5" --pidfile "${TD_AGENT_PID_FILE}" --name ruby || RETVAL="$?"
  [ "$RETVAL" = 2 ] && return 2
  # Many daemons don't delete their pidfiles when they exit.
  rm -f "${TD_AGENT_PID_FILE}"
  return "$RETVAL"
}

#
# Function that sends a SIGHUP to the daemon/service
#
do_reload() {
  #
  # If the daemon can reload its configuration without
  # restarting (for example, when it is sent a SIGHUP),
  # then implement that here.
  #
  start-stop-daemon --stop --signal 1 --quiet --pidfile "${TD_AGENT_PID_FILE}" --name ruby
}

do_configtest() {
  eval "${TD_AGENT_ARGS} --user ${TD_AGENT_USER} --group ${TD_AGENT_GROUP} --dry-run -q"
}

RETVAL=0
case "$1" in
"start" )
  log_daemon_msg "Starting ${TD_AGENT_NAME} " "${TD_AGENT_NAME}"
  do_start || RETVAL="$?"
  case "$RETVAL" in
  0 | 1 ) log_end_msg 0 ;;
      2 ) log_end_msg 1 ;;
  esac
  ;;
"stop" )
  log_daemon_msg "Stopping ${TD_AGENT_NAME} " "${TD_AGENT_NAME}"
  do_stop || RETVAL="$?"
  case "$RETVAL" in
  0 | 1 ) log_end_msg 0 ;;
      2 ) log_end_msg 1 ;;
  esac
  ;;
"reload" | "force-reload" )
  #
  # If do_reload() is not implemented then leave this commented out
  # and leave 'force-reload' as an alias for 'restart'.
  #
  log_daemon_msg "Reloading ${TD_AGENT_NAME} " "${TD_AGENT_NAME}"
  do_configtest || log_end_msg 1
  do_reload || RETVAL="$?"
  log_end_msg "$RETVAL"
  ;;
"restart" )
  #
  # If the "reload" option is implemented then remove the
  # 'force-reload' alias
  #
  log_daemon_msg "Restarting ${TD_AGENT_NAME} " "${TD_AGENT_NAME}"
  do_configtest || log_end_msg 1
  do_stop || RETVAL="$?"
  case "$RETVAL" in
  0 | 1 )
    RETVAL=0
    do_start || RETVAL="$?"
    case "$RETVAL" in
    0 ) log_end_msg 0 ;;
    1 ) log_end_msg 1 ;; # Old process is still running
    * ) log_end_msg 1 ;; # Failed to start
    esac
    ;;
  * )
    # Failed to stop
    log_end_msg 1
    ;;
  esac
  ;;
"status" )
  status_of_proc "${TD_AGENT_RUBY}" "${TD_AGENT_NAME}"
  ;;
"configtest" )
  do_configtest || RETVAL="$?"
  log_end_msg "$RETVAL"
  ;;
* )
  echo "Usage: $0 {start|stop|reload|restart|force-reload|status|configtest}" >&2
  exit 1
  ;;
esac
