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


# The Hadoop command script
#
# Environment Variables
#
#   JAVA_HOME        The java implementation to use.  Overrides JAVA_HOME.
#
#   YARN_USER_CLASSPATH Additional user CLASSPATH entries.
#
#   YARN_USER_CLASSPATH_FIRST  If set to non empty value then the user classpath
#                              specified in YARN_USER_CLASSPATH will be
#                              appended at the beginning of YARN's final
#                              classpath instead of at the end.
#
#   YARN_HEAPSIZE  The maximum amount of heap to use, in MB. 
#                    Default is 1000.
#
#   YARN_{COMMAND}_HEAPSIZE overrides YARN_HEAPSIZE for a given command
#                           eg YARN_NODEMANAGER_HEAPSIZE sets the heap
#                           size for the NodeManager.  If you set the
#                           heap size in YARN_{COMMAND}_OPTS or YARN_OPTS
#                           they take precedence.
#
#   YARN_OPTS      Extra Java runtime options.
#   
#   YARN_CLIENT_OPTS         when the respective command is run.
#   YARN_{COMMAND}_OPTS etc  YARN_NODEMANAGER_OPTS applies to NodeManager 
#                              for e.g.  YARN_CLIENT_OPTS applies to 
#                              more than one command (fs, dfs, fsck, 
#                              dfsadmin etc)  
#
#   YARN_CONF_DIR  Alternate conf dir. Default is ${HADOOP_YARN_HOME}/conf.
#
#   YARN_ROOT_LOGGER The root appender. Default is INFO,console
#

bin=`dirname "${BASH_SOURCE-$0}"`
bin=`cd "$bin"; pwd`

DEFAULT_LIBEXEC_DIR="$bin"/../libexec

HADOOP_LIBEXEC_DIR=${HADOOP_LIBEXEC_DIR:-$DEFAULT_LIBEXEC_DIR}
. $HADOOP_LIBEXEC_DIR/yarn-config.sh

function print_usage(){
  echo "Usage: yarn [--config confdir] [COMMAND | CLASSNAME]"
  echo "  CLASSNAME                             run the class named CLASSNAME"
  echo " or"
  echo "  where COMMAND is one of:"
  echo "  resourcemanager                       run the ResourceManager"
  echo "                                        Use -format-state-store for deleting the RMStateStore."
  echo "                                        Use -remove-application-from-state-store <appId> for "
  echo "                                            removing application from RMStateStore."
  echo "  nodemanager                           run a nodemanager on each slave"
  echo "  timelineserver                        run the timeline server"
  echo "  rmadmin                               admin tools"
  echo "  router                                run the Router daemon"
  echo "  sharedcachemanager                    run the SharedCacheManager daemon"
  echo "  scmadmin                              SharedCacheManager admin tools"
  echo "  version                               print the version"
  echo "  jar <jar>                             run a jar file"
  echo "  application                           prints application(s)"
  echo "                                        report/kill application"
  echo "  applicationattempt                    prints applicationattempt(s)"
  echo "                                        report"
  echo "  container                             prints container(s) report"
  echo "  node                                  prints node report(s)"
  echo "  queue                                 prints queue information"
  echo "  logs                                  dump container logs"
  echo "  schedulerconf                         updates scheduler configuration"
  echo "  classpath                             prints the class path needed to"
  echo "                                        get the Hadoop jar and the"
  echo "                                        required libraries"
  echo "  cluster                               prints cluster information"
  echo "  daemonlog                             get/set the log level for each"
  echo "                                        daemon"
  echo "  top                                   run cluster usage tool"
  echo ""
  echo "Most commands print help when invoked w/o parameters."
}

# if no args specified, show usage
if [ $# = 0 ]; then
  print_usage
  exit 1
fi

# get arguments
COMMAND=$1
shift

case $COMMAND in
  # usage flags
  --help|-help|-h)
    print_usage
    exit
    ;;
esac

if [ -f "${YARN_CONF_DIR}/yarn-env.sh" ]; then
  . "${YARN_CONF_DIR}/yarn-env.sh"
fi

# some Java parameters
if [ "$JAVA_HOME" != "" ]; then
  #echo "run java in $JAVA_HOME"
  JAVA_HOME=$JAVA_HOME
fi
  
if [ "$JAVA_HOME" = "" ]; then
  echo "Error: JAVA_HOME is not set."
  exit 1
fi

JAVA=$JAVA_HOME/bin/java
JAVA_HEAP_MAX=-Xmx1000m 

# check envvars which might override default args
if [ "$YARN_HEAPSIZE" != "" ]; then
  #echo "run with heapsize $YARN_HEAPSIZE"
  JAVA_HEAP_MAX="-Xmx""$YARN_HEAPSIZE""m"
  #echo $JAVA_HEAP_MAX
fi

# CLASSPATH initially contains $HADOOP_CONF_DIR & $YARN_CONF_DIR
if [ ! -d "$HADOOP_CONF_DIR" ]; then
  echo No HADOOP_CONF_DIR set. 
  echo Please specify it either in yarn-env.sh or in the environment.
  exit 1
fi

CLASSPATH="${HADOOP_CONF_DIR}:${YARN_CONF_DIR}:${CLASSPATH}"

# for developers, add Hadoop classes to CLASSPATH
if [ -d "$HADOOP_YARN_HOME/yarn-api/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-api/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-common/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-common/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-mapreduce/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-mapreduce/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-master-worker/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-master-worker/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-server/yarn-server-nodemanager/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-nodemanager/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-server/yarn-server-common/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-common/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-server/yarn-server-resourcemanager/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-resourcemanager/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/yarn-server/yarn-server-applicationhistoryservice/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-applicationhistoryservice/target/classes
fi
if [ -d "$HADOOP_YARN_HOME/build/test/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/target/test/classes
fi
if [ -d "$HADOOP_YARN_HOME/build/tools" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/build/tools
fi
if [ -d "$HADOOP_YARN_HOME/yarn-server/yarn-server-router/target/classes" ]; then
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-router/target/classes
fi

CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/${YARN_DIR}/*
CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/${YARN_LIB_JARS_DIR}/*

# Add user defined YARN_USER_CLASSPATH to the class path (if defined)
if [ -n "$YARN_USER_CLASSPATH" ]; then
  if [ -n "$YARN_USER_CLASSPATH_FIRST" ]; then
    # User requested to add the custom entries at the beginning
    CLASSPATH=${YARN_USER_CLASSPATH}:${CLASSPATH}
  else
    # By default we will just append the extra entries at the end
    CLASSPATH=${CLASSPATH}:${YARN_USER_CLASSPATH}
  fi
fi

# so that filenames w/ spaces are handled correctly in loops below
IFS=

# default log directory & file
if [ "$YARN_LOG_DIR" = "" ]; then
  YARN_LOG_DIR="$HADOOP_YARN_HOME/logs"
fi
if [ "$YARN_LOGFILE" = "" ]; then
  YARN_LOGFILE='yarn.log'
fi

# restore ordinary behaviour
unset IFS

# figure out which class to run
if [ "$COMMAND" = "classpath" ] ; then
  if [ "$#" -gt 0 ]; then
    CLASS=org.apache.hadoop.util.Classpath
  else
    if [ "$cygwin" = true ]; then
      CLASSPATH=$(cygpath -p -w "$CLASSPATH" 2>/dev/null)
    fi
    echo $CLASSPATH
    exit 0
  fi
elif [ "$COMMAND" = "rmadmin" ] ; then
  CLASS='org.apache.hadoop.yarn.client.cli.RMAdminCLI'
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "router" ] ; then
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/rt-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.router.Router'
  YARN_OPTS="$YARN_OPTS $YARN_ROUTER_OPTS"
elif [ "$COMMAND" = "scmadmin" ] ; then
  CLASS='org.apache.hadoop.yarn.client.SCMAdmin'
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "application" ] || 
     [ "$COMMAND" = "applicationattempt" ] || 
     [ "$COMMAND" = "container" ]; then
  CLASS=org.apache.hadoop.yarn.client.cli.ApplicationCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
  set -- $COMMAND $@
elif [ "$COMMAND" = "node" ] ; then
  CLASS=org.apache.hadoop.yarn.client.cli.NodeCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "queue" ] ; then
  CLASS=org.apache.hadoop.yarn.client.cli.QueueCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "resourcemanager" ] ; then
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/rm-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.resourcemanager.ResourceManager'
  YARN_OPTS="$YARN_OPTS $YARN_RESOURCEMANAGER_OPTS"
  if [ "$YARN_RESOURCEMANAGER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_RESOURCEMANAGER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "historyserver" ] ; then
  echo "DEPRECATED: Use of this command to start the timeline server is deprecated." 1>&2
  echo "Instead use the timelineserver command for it." 1>&2
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/ahs-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.applicationhistoryservice.ApplicationHistoryServer'
  YARN_OPTS="$YARN_OPTS $YARN_HISTORYSERVER_OPTS"
  if [ "$YARN_HISTORYSERVER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_HISTORYSERVER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "timelineserver" ] ; then
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/timelineserver-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.applicationhistoryservice.ApplicationHistoryServer'
  YARN_OPTS="$YARN_OPTS $YARN_TIMELINESERVER_OPTS"
  if [ "$YARN_TIMELINESERVER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_TIMELINESERVER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "sharedcachemanager" ] ; then
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/scm-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.sharedcachemanager.SharedCacheManager'
  YARN_OPTS="$YARN_OPTS $YARN_SHAREDCACHEMANAGER_OPTS"
  if [ "$YARN_SHAREDCACHEMANAGER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_SHAREDCACHEMANAGER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "nodemanager" ] ; then
  CLASSPATH=${CLASSPATH}:$YARN_CONF_DIR/nm-config/log4j.properties
  CLASS='org.apache.hadoop.yarn.server.nodemanager.NodeManager'
  YARN_OPTS="$YARN_OPTS -server $YARN_NODEMANAGER_OPTS"
  if [ "$YARN_NODEMANAGER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_NODEMANAGER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "proxyserver" ] ; then
  CLASS='org.apache.hadoop.yarn.server.webproxy.WebAppProxyServer'
  YARN_OPTS="$YARN_OPTS $YARN_PROXYSERVER_OPTS"
  if [ "$YARN_PROXYSERVER_HEAPSIZE" != "" ]; then
    JAVA_HEAP_MAX="-Xmx""$YARN_PROXYSERVER_HEAPSIZE""m"
  fi
elif [ "$COMMAND" = "version" ] ; then
  CLASS=org.apache.hadoop.util.VersionInfo
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "jar" ] ; then
  CLASS=org.apache.hadoop.util.RunJar
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "logs" ] ; then
  CLASS=org.apache.hadoop.yarn.client.cli.LogsCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "daemonlog" ] ; then
  CLASS=org.apache.hadoop.log.LogLevel
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "cluster" ] ; then
  CLASS=org.apache.hadoop.yarn.client.cli.ClusterCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "schedulerconf" ] ; then
  CLASS=org.apache.hadoop.yarn.client.cli.SchedConfCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
elif [ "$COMMAND" = "top" ]; then
  CLASS=org.apache.hadoop.yarn.client.cli.TopCLI
  YARN_OPTS="$YARN_OPTS $YARN_CLIENT_OPTS"
  doNotSetCols=0
  doNotSetRows=0
  for i in "$@"; do
    if [[ $i == "-cols" ]]; then
      doNotSetCols=1
    fi
    if [[ $i == "-rows" ]]; then
      doNotSetRows=1
    fi
  done
  if [ $doNotSetCols == 0 ] && [ -n "${TERM}" ]; then
    cols=$(tput cols)
    if [ -n "$cols" ]; then
      args=( $@ )
      args=("${args[@]}" "-cols" "$cols")
      set -- "${args[@]}"
    fi
  fi
  if [ $doNotSetRows == 0 ] && [ -n "${TERM}" ]; then
    rows=$(tput lines)
    if [ -n "$rows" ]; then
      args=( $@ )
      args=("${args[@]}" "-rows" "$rows")
      set -- "${args[@]}"
    fi
  fi
else
  CLASS=$COMMAND
fi

# cygwin path translation
if [ "$cygwin" = true ]; then
  CLASSPATH=$(cygpath -p -w "$CLASSPATH" 2>/dev/null)
  HADOOP_LOG_DIR=$(cygpath -w "$HADOOP_LOG_DIR" 2>/dev/null)
  HADOOP_PREFIX=$(cygpath -w "$HADOOP_PREFIX" 2>/dev/null)
  HADOOP_CONF_DIR=$(cygpath -w "$HADOOP_CONF_DIR" 2>/dev/null)
  HADOOP_COMMON_HOME=$(cygpath -w "$HADOOP_COMMON_HOME" 2>/dev/null)
  HADOOP_HDFS_HOME=$(cygpath -w "$HADOOP_HDFS_HOME" 2>/dev/null)
  HADOOP_YARN_HOME=$(cygpath -w "$HADOOP_YARN_HOME" 2>/dev/null)
  HADOOP_MAPRED_HOME=$(cygpath -w "$HADOOP_MAPRED_HOME" 2>/dev/null)
  YARN_LOG_DIR=$(cygpath -w "$YARN_LOG_DIR" 2>/dev/null)
fi

YARN_OPTS="$YARN_OPTS -Dhadoop.log.dir=$YARN_LOG_DIR"
YARN_OPTS="$YARN_OPTS -Dyarn.log.dir=$YARN_LOG_DIR"
YARN_OPTS="$YARN_OPTS -Dhadoop.log.file=$YARN_LOGFILE"
YARN_OPTS="$YARN_OPTS -Dyarn.log.file=$YARN_LOGFILE"
YARN_OPTS="$YARN_OPTS -Dyarn.home.dir=$HADOOP_YARN_HOME"

if [ "$cygwin" = true ]; then
  HADOOP_HOME=$(cygpath -w "$HADOOP_PREFIX" 2>/dev/null)
  YARN_OPTS="$YARN_OPTS -Dhadoop.home.dir=$HADOOP_HOME"
  export HADOOP_HOME
else
  YARN_OPTS="$YARN_OPTS -Dhadoop.home.dir=$HADOOP_PREFIX"
fi

YARN_OPTS="$YARN_OPTS -Dhadoop.root.logger=${YARN_ROOT_LOGGER:-${HADOOP_LOGLEVEL},console}"
YARN_OPTS="$YARN_OPTS -Dyarn.root.logger=${YARN_ROOT_LOGGER:-${HADOOP_LOGLEVEL},console}"
if [ "x$JAVA_LIBRARY_PATH" != "x" ]; then
  if [ "$cygwin" = true ]; then
    JAVA_LIBRARY_PATH=$(cygpath -w "$JAVA_LIBRARY_PATH" 2>/dev/null)
  fi
  YARN_OPTS="$YARN_OPTS -Djava.library.path=$JAVA_LIBRARY_PATH"
fi  

exec "$JAVA" -Dproc_$COMMAND $JAVA_HEAP_MAX $YARN_OPTS -classpath "$CLASSPATH" $CLASS "$@"