#!/bin/bash

# Copyright (c) 2012 Cloudera, Inc. All rights reserved.

# Time marker for both stderr and stdout
date; date 1>&2

cloudera_config=`dirname $0`
cloudera_config=`cd "$cloudera_config/../common"; pwd`
. ${cloudera_config}/cloudera-config.sh

# Load the parcel environment
source_parcel_environment

# attempt to find java
locate_cdh_java_home

# simulate /etc/default/hadoop
. ${cloudera_config}/cdh-default-hadoop

echo "using $JAVA_HOME as JAVA_HOME"
echo "using $CDH_VERSION as CDH_VERSION"
echo "using $CDH_YARN_HOME as CDH_YARN_HOME"
echo "using $CDH_MR2_HOME as CDH_MR2_HOME"
echo "using $CONF_DIR as CONF_DIR"

export YARN_CONF_DIR=$CONF_DIR

YARN=$CDH_YARN_HOME/bin/yarn

set_hadoop_classpath

# debug
set -x

# Search-replace {{CMF_CONF_DIR}} in files
replace_conf_dir

# Update path to cgroups prefix in the yarn config
perl -pi -e "s#{{CGROUP_GROUP_CPU}}#$CGROUP_GROUP_CPU#g" $CONF_DIR/yarn-site.xml

# replace {{PID}} in HeapDumpPath option
export YARN_NODEMANAGER_OPTS=$(replace_pid $YARN_NODEMANAGER_OPTS)
export YARN_RESOURCEMANAGER_OPTS=$(replace_pid $YARN_RESOURCEMANAGER_OPTS)
export HADOOP_JOB_HISTORYSERVER_OPTS=$(replace_pid $HADOOP_JOB_HISTORYSERVER_OPTS)

# replace {{CONF_DIR}} placeholder with real $CONF_DIR
export YARN_NODEMANAGER_OPTS=$(replace_conf_dir_env $YARN_NODEMANAGER_OPTS)
export YARN_RESOURCEMANAGER_OPTS=$(replace_conf_dir_env $YARN_RESOURCEMANAGER_OPTS)
export HADOOP_JOB_HISTORYSERVER_OPTS=$(replace_conf_dir_env $HADOOP_JOB_HISTORYSERVER_OPTS)

# Make sure topology.py is executable
make_scripts_executable

# Disable IPv6.
export YARN_OPTS="-Djava.net.preferIPv4Stack=true $YARN_OPTS"

acquire_kerberos_tgt yarn.keytab

if [ $1 = "RefreshQueuesAndNodes" ]; then
  # If RM is in HA mode, we need to issue command only to active RM
  if [ -n $2 ] && [ $2 = "true" ]; then
    # find if RM is in standby mode
    HASTATE="$($YARN --config $CONF_DIR rmadmin -getServiceState $3)"
    if [[ $HASTATE =~ "standby" ]]; then
      # Standby RM have no state to refresh, exit successfully
      cm_success_and_exit YARN_RM_IN_STANDBY_MODE
    fi
  fi

  $YARN --config $CONF_DIR rmadmin -refreshQueues
  RET=$?
  if [ $RET -eq 0 ]; then
    echo "RefreshQueues completed successfully"
    if [ -n $4 ] && [ $4 != "unused" ]; then
      $YARN --config $CONF_DIR rmadmin -refreshNodes -g $4 -client
    else
      $YARN --config $CONF_DIR rmadmin -refreshNodes
    fi
    RET=$?
    if [ $RET -eq 0 ]; then
      echo "RefreshNodes completed successfully"
    else
      echo "RefreshNodes failed"
    fi
  else
    echo "RefreshQueues failed. Skipping RefreshNodes."
  fi
  exit $RET
elif [ $1 = "historyserver" ]; then
  # historyserver is actually a MR2 daemon.
  MAPRED=$HADOOP_MAPRED_HOME/bin/mapred
  exec $MAPRED "$@"
elif [ $1 = "application-diagnostic-data" ]; then
  set -e
  PYTHON_CMD=$CMF_PACKAGE_DIR/../agent/build/env/bin/python
  COLLECT_APPLICATION_DATA="$(dirname $0)/../support/collect_jobs_stats/collect_jobs_stats.py"
  PROCESS_DIR_NAME=$(basename $(pwd))
  DIAGNOSTICS_DUMP_DIR="$HEAP_DUMP_DIR/$PROCESS_DIR_NAME"
  COLLECT_APP_DATA_ARGS="--app_ids $3 --output_dir $DIAGNOSTICS_DUMP_DIR/diagnostics --hadoop_conf_dir $CONF_DIR"
  $PYTHON_CMD $COLLECT_APPLICATION_DATA $COLLECT_APP_DATA_ARGS
  tar -czf $DIAGNOSTICS_DUMP_DIR.tar.gz -C $DIAGNOSTICS_DUMP_DIR diagnostics
  rm -rf $DIAGNOSTICS_DUMP_DIR
  ln -s $DIAGNOSTICS_DUMP_DIR.tar.gz output.gz
elif [ $1 = "install-mr-framework" ]; then
  acquire_kerberos_tgt yarn.keytab
  # Expected args:
  # $2 = dir in which mr-framework.tar.gz will be copied.
  # $3 = name of file
  MR_FRAMEWORK_DIR=$2
  MR_FRAMEWORK_FILENAME=$3
  UNPACKED_DIR=$(mktemp -d -t mapreduce-framework-staging.XXXXXXXXXX)
  tar -cvzhf $UNPACKED_DIR/$MR_FRAMEWORK_FILENAME --directory=$HADOOP_MAPRED_HOME .
  $CDH_HADOOP_BIN --config $CONF_DIR fs -rm -skipTrash $MR_FRAMEWORK_DIR/$MR_FRAMEWORK_FILENAME
  exec $CDH_HADOOP_BIN --config $CONF_DIR fs -put $UNPACKED_DIR/$MR_FRAMEWORK_FILENAME $MR_FRAMEWORK_DIR/
else
  if [ $1 = "nodemanager" ]; then
    # In secure mode, container are launched as user who submitted the job and
    # needs read access to the config files. We don't make everything world-readable
    # because this directory also contains files critical from security standpoint.
    chmod go+r $CONF_DIR/yarn-site.xml
    chmod go+r $CONF_DIR/hdfs-site.xml
    chmod go+r $CONF_DIR/core-site.xml
    chmod go+r $CONF_DIR/mapred-site.xml
    chmod go+r $CONF_DIR/ssl-client.xml
    chmod go+r $CONF_DIR/hadoop-policy.xml

    # Below logic is a quick fix to OPSAPS-28378, see the jira for details.
    if [ "$IS_KERBERIZED" = "true" ]; then
      for nm_local_dir in "$NM_LOCAL_DIRS"; do
        for dir in "$nm_local_dir"/usercache/*; do
          if [ -d "$dir" ]; then
            base=$(basename $dir)
            owner=$(stat -c %U $dir)
            if [ "$owner" != "$base" ]; then
              echo "directory owner doesn't match directory name, will try to chown"
              if chown -R $base $dir ; then
                echo "chown succeeded"
              else
                echo "chown failed, will try to remove the directory"
                if rm -rf "$dir" ; then
                  echo "rm succeeded"
                else
                  echo "rm failed"
                fi
              fi
            fi
          fi
        done
      done
    fi
  fi
  exec $YARN "$@"
fi
