#!/bin/sh
### BEGIN INIT INFO
# Provides:          td-agent
# Required-Start:    $network $local_fs
# Required-Stop:     $network $local_fs
# Default-Start:     2 3 4 5
# Default-Stop:      0 1 6
# Short-Description: data collector for Treasure Data
# Description:       td-agent is a data collector
### END INIT INFO
# pidfile:           /var/run/td-agent/td-agent.pid

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

kill_by_file() {
  local sig="$1"
  shift 1
  local pid="$(cat "$@" 2>/dev/null || true)"
  if [ -n "${pid}" ]; then
    if /bin/kill "${sig}" "${pid}" 1>/dev/null 2>&1; then
      return 0
    else
      return 2
    fi
  else
    return 1
  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
  if kill_by_file -0 "${TD_AGENT_PID_FILE}"; then
    return 1
  else
    start-stop-daemon --start --quiet --pidfile "${TD_AGENT_PID_FILE}" --exec "${TD_AGENT_RUBY}" \
      ${START_STOP_DAEMON_ARGS} -- ${TD_AGENT_ARGS} \
      || return 2
  fi
  # 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
  if [ -e "${TD_AGENT_PID_FILE}" ]; then
    # Use own process termination instead of killproc because killproc can't wait SIGTERM
    if kill_by_file -TERM "${TD_AGENT_PID_FILE}"; then
      local i
      for i in $(seq "${STOPTIMEOUT}"); do
        if kill_by_file -0 "${TD_AGENT_PID_FILE}"; then
          sleep 1
        else
          break
        fi
      done
      if kill_by_file -0 "${TD_AGENT_PID_FILE}"; then
        echo -n "Timeout error occurred trying to stop ${TD_AGENT_NAME}..."
        return 2
      else
        rm -f "${TD_AGENT_PID_FILE}"
      fi
    else
      return 1
    fi
  else
    if killproc "${TD_AGENT_PROG_NAME:-${TD_AGENT_NAME}}"; then
      rm -f "${TD_AGENT_PID_FILE}"
    else
      return 2
    fi
  fi
}

#
# Function that sends a SIGHUP to the daemon/service
#
do_reload() {
  kill_by_file -HUP "${TD_AGENT_PID_FILE}"
}

do_restart() {
  if ! do_configtest; then
    return 1
  fi
  local val=0
  do_stop || val="$?"
  case "${val}" in
  0 | 1 )
    if ! do_start; then
      return 1
    fi
    ;;
  * ) # Failed to stop
    return 1
    ;;
  esac
}

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

RETVAL=0
case "$1" in
"start" )
  echo -n "Starting ${TD_AGENT_NAME}: "
  do_start || RETVAL="$?"
  case "$RETVAL" in
  0 | 1 )
    log_success_msg "${TD_AGENT_NAME}"
    ;;
  * )
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
    ;;
  esac
  ;;
"stop" )
  echo -n "Stopping ${TD_AGENT_NAME}: "
  do_stop || RETVAL="$?"
  case "$RETVAL" in
  0 )
    log_success_msg "${TD_AGENT_NAME}"
    ;;
  * )
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
    ;;
  esac
  ;;
"reload" | "force-reload" )
  echo -n "Reloading ${TD_AGENT_NAME}: "
  if ! do_configtest; then
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
  fi
  if do_reload; then
    log_success_msg "${TD_AGENT_NAME}"
  else
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
  fi
  ;;
"restart" )
  echo -n "Restarting ${TD_AGENT_NAME}: "
  if do_restart; then
    log_success_msg "${TD_AGENT_NAME}"
  else
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
  fi
  ;;
"status" )
  if kill_by_file -0 "${TD_AGENT_PID_FILE}"; then
    log_success_msg "${TD_AGENT_NAME} is running"
  else
    log_failure_msg "${TD_AGENT_NAME} is not running"
    exit 1
  fi
  ;;
"configtest" )
  if do_configtest; then
    log_success_msg "${TD_AGENT_NAME}"
  else
    log_failure_msg "${TD_AGENT_NAME}"
    exit 1
  fi
  ;;
* )
  echo "Usage: $0 {start|stop|reload|restart|force-reload|status|configtest}" >&2
  exit 1
  ;;
esac
