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

bin_dir=$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )
at_home=$( cd "$( dirname "$bin_dir" )" && pwd )

function print_usage() {
  cat <<EOF

Usage: agitator <command>

Possible commands:
  start     Starts agitator
  stop      Stop agitator
EOF
}

if [ -f "$at_home/conf/env.sh" ]; then
  . "$at_home"/conf/env.sh
else
  . "$at_home"/conf/env.sh.example
fi

function start_agitator() {
  hash pssh 2>/dev/null || { echo >&2 "The agitator requires pssh to be installed. Aborting."; exit 1; }

  mkdir -p "${at_home}/logs"
  log_base="${at_home}/logs/$(date +%Y%m%d%H%M%S)_$(hostname)"
  libexec="${at_home}/libexec"
  master_log="${log_base}_master-agitator"
  tserver_log="${log_base}_tserver-agitator"
  datanode_log="${log_base}_datanode-agitator"
  master_cmd="nohup ${libexec}/master-agitator.pl $AGTR_MASTER_KILL_SLEEP_TIME $AGTR_MASTER_RESTART_SLEEP_TIME"
  tserver_cmd="nohup ${libexec}/tserver-agitator.pl $AGTR_TSERVER_KILL_SLEEP_TIME $AGTR_TSERVER_RESTART_SLEEP_TIME $AGTR_TSERVER_MIN_KILL $AGTR_TSERVER_MAX_KILL"
  datanode_cmd="nohup ${libexec}/datanode-agitator.pl $AGTR_DATANODE_KILL_SLEEP_TIME $AGTR_DATANODE_RESTART_SLEEP_TIME $HADOOP_HOME $AGTR_DATANODE_MIN_KILL $AGTR_DATANODE_MAX_KILL"
  [[ -n $AGITATOR_USER ]] || AGITATOR_USER=$(whoami)

  if [[ $AGITATOR_USER == root ]];  then
    echo "Running master-agitator and tserver-agitator as $AGTR_ACCUMULO_USER using su. Running datanode-agitator as $AGTR_HDFS_USER using su."
    su -c "$master_cmd >${master_log}.out 2>${master_log}.err" & -m - "$AGTR_ACCUMULO_USER"
    su -c "$tserver_cmd >${tserver_log}.out 2>${tserver_log}.err" & -m - "$AGTR_ACCUMULO_USER"
    su -c "$datanode_cmd >${datanode_log}.out 2>${datanode_log}.err" & -m - "$AGTR_HDFS_USER"
  else
    if [[ $AGITATOR_USER == "$AGTR_ACCUMULO_USER" ]]; then
      echo "Running master-agitator and tserver-agitator as $AGITATOR_USER"
      $master_cmd > "${master_log}.out" 2> "${master_log}.err" &
      $tserver_cmd > "${tserver_log}.out" 2> "${tserver_log}.err" &
    else
      echo "Running master-agitator and tserver-agitator as $AGTR_ACCUMULO_USER using sudo."
      sudo -u "$AGTR_ACCUMULO_USER" $master_cmd > "${master_log}.out" 2> "${master_log}.err" &
      sudo -u "$AGTR_ACCUMULO_USER" $tserver_cmd > "${tserver_log}.out" 2> "${tserver_log}.err" &
    fi
    if [[ $AGITATOR_USER == "$AGTR_HDFS_USER" ]]; then
      echo "Running datanode-agitator as $AGITATOR_USER"
      $datanode_cmd > "${datanode_log}.out" 2> "${datanode_log}.err" &
    else
      echo "Running datanode-agitator as $AGTR_HDFS_USER using sudo."
      sudo -u "$AGTR_HDFS_USER" $datanode_cmd > "${datanode_log}.out" 2> "${datanode_log}.err" &
    fi
  fi

  if ${AGTR_HDFS:-false} ; then
    agitator_log=${log_base}_hdfs-agitator
    sudo -u "$AGTR_HDFS_SUPERUSER" nohup "${libexec}/hdfs-agitator.pl" --sleep "${AGTR_HDFS_SLEEP_TIME}" --hdfs-cmd "${AGTR_HDFS_COMMAND}" --superuser "${AGTR_HDFS_SUPERUSER}" >"${agitator_log}.out" 2>"${agitator_log}.err" &
  fi
}

function stop_agitator() {
  [[ -n $AGITATOR_USER ]] || AGITATOR_USER=$(whoami)
  if [[ $AGITATOR_USER == root ]]; then
    echo "Stopping all processes matching 'agitator.pl' as root"
    pkill -f agitator.pl 2>/dev/null
  elif [[ $AGITATOR_USER == "$AGTR_ACCUMULO_USER" ]];  then
    echo "Stopping all processes matching 'datanode-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f datanode-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'hdfs-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f hdfs-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'agitator.pl' as $AGITATOR_USER"
    pkill -f agitator.pl 2>/dev/null 2>/dev/null
  else
    echo "Stopping all processes matching 'datanode-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f datanode-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'hdfs-agitator.pl' as $AGTR_HDFS_USER"
    sudo -u "$AGTR_HDFS_USER" pkill -f hdfs-agitator.pl 2>/dev/null
    echo "Stopping all processes matching 'agitator.pl' as $AGTR_ACCUMULO_USER"
    sudo -u "$AGTR_ACCUMULO_USER" pkill -f agitator.pl 2>/dev/null
  fi
}

case "$1" in
  start)
    start_agitator
    ;;
  stop)
    stop_agitator
    ;;
  *)
    echo "ERROR: unknown command - $2"
    print_usage
    exit 1
esac