#!/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" > /dev/null; 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"
  echo "where COMMAND is one of:"
  echo "  resourcemanager      run the ResourceManager" 
  echo "  nodemanager          run a nodemanager on each slave" 
  echo "  historyserver        run the application history server" 
  echo "  rmadmin              admin tools" 
  echo "  version              print the version"
  echo "  jar <jar>            run a jar file"
  echo "  application          prints application(s) report/kill application"
  echo "  applicationattempt   prints applicationattempt(s) report"
  echo "  container            prints container(s) report"
  echo "  node                 prints node report(s)"
  echo "  logs                 dump container logs"
  echo "  classpath            prints the class path needed to get the"
  echo "                       Hadoop jar and the required libraries"
  echo "  daemonlog            get/set the log level for each daemon"
  echo " or"
  echo "  CLASSNAME            run the class named CLASSNAME"
  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
  CLASSPATH=${CLASSPATH}:$HADOOP_YARN_HOME/yarn-server/yarn-server-applicationhistoryservice/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

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
  echo $CLASSPATH
  exit
elif [ "$COMMAND" = "rmadmin" ] ; then
  CLASS='org.apache.hadoop.yarn.client.cli.RMAdminCLI'
  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" = "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
  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" = "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"
else
  CLASS=$COMMAND
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"
YARN_OPTS="$YARN_OPTS -Dhadoop.home.dir=$HADOOP_YARN_HOME"
YARN_OPTS="$YARN_OPTS -Dhadoop.root.logger=${YARN_ROOT_LOGGER:-INFO,console}"
YARN_OPTS="$YARN_OPTS -Dyarn.root.logger=${YARN_ROOT_LOGGER:-INFO,console}"
if [ "x$JAVA_LIBRARY_PATH" != "x" ]; then
  YARN_OPTS="$YARN_OPTS -Djava.library.path=$JAVA_LIBRARY_PATH"
fi  

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