#!/bin/sh
#
# Copyright (C) 2003-2012 Intel Corporation.  All Rights Reserved.
# 
# The source code contained or described herein and all documents
# related to the source code ("Material") are owned by Intel Corporation
# or its suppliers or licensors.  Title to the Material remains with
# Intel Corporation or its suppliers and licensors.  The Material is
# protected by worldwide copyright and trade secret laws and treaty
# provisions.  No part of the Material may be used, copied, reproduced,
# modified, published, uploaded, posted, transmitted, distributed, or
# disclosed in any way without Intel's prior express written permission.
# 
# No license under any patent, copyright, trade secret or other
# intellectual property right is granted to or conferred upon you by
# disclosure or delivery of the Materials, either expressly, by
# implication, inducement, estoppel or otherwise.  Any license under
# such intellectual property rights must be express and approved by
# Intel in writing.
#

# This utility is for cleaning up nodes

base_name=$(basename $0)
Usage()
{
    if [ $silent -eq 0 ]; then
        echo "Usage: $(basename $0) [-i <input_file> | -t -f <hostsfile>] [-r <rshcmd>] [-b <branch_count>] \\"
        echo "                  [-p] [-s | -d] [-h] [-V]"
        echo "   or"
        echo "       $(basename $0) [--input <input_file> | --total --file <hostsfile>] [--rsh <rshcmd>] [--branch <branch_count>] \\"
        echo "                  [--parallel] [--silent | --verbose] [--help] [--version]"
        echo "  Options:"
        echo "      -i file_name        specify input file name"
        echo "      -t                  remove all user's processes totally"
        echo "      -f hostfile         set host file to clean"
        echo "      -r rsh              specify remote shell (ssh by default)"
        echo "      -b branch_count     set branch count to use"
        echo "      -p                  use parallel mode"
        echo "      -s                  suppress the stdout/err info, silent mode"
        echo "      -d                  verbose info"
        echo "      -h                  prints this usage"
        echo "      -V                  version of library"
        echo ""
        echo "Intel(R) MPI Library for Linux*, Version 4.1 Build 20120831"
        echo "Copyright (C) 2012 Intel Corporation. All rights reserved"
    fi
    exit 0 
}

print_version()
{
    if [ $silent -eq 0 ]; then
        echo "Intel(R) MPI Library for Linux*, Version 4.1"
        echo "Build 20120831 Platform Intel(R) 64 64-bit applications"
        echo "Copyright (C) 2012 Intel Corporation. All rights reserved"
    fi
    exit 0
}

# Default settings
rsh_cmd="ssh -x -q"
username=""
input_file=""
host_file=""
branch_count=32
parallel_mode=0
total_mode=0
verbose=0
silent=0
host_name=`hostname -s`

opts_to_propagate=""

# Parse arguments
while [ $# -gt 0 ]; do
    arg="$1"
    if [ -n "$arg" ]; then
        case "$arg" in
            -i | --input )      input_file=$2; shift ;;
            -t | --total )      total_mode=1 ; opts_to_propagate="${opts_to_propagate} $1" ;;
            -f | --file )       host_file=$2; shift ;;
            -r | --rsh )        rsh_cmd=$2; opts_to_propagate="${opts_to_propagate} $1 $2"; shift ;;
            -b | --branch )     branch_count=$2; opts_to_propagate="${opts_to_propagate} $1 $2"; shift ;;
            -p | --parallel )   parallel_mode=1 ; opts_to_propagate="${opts_to_propagate} $1" ;;
            -s | --silent )      silent=1; verbose=0 ; opts_to_propagate="${opts_to_propagate} $1" ;;
            -d | --verbose )    verbose=1; silent=0 ; opts_to_propagate="${opts_to_propagate} $1" ;;
            -h | --help )       Usage ;;
            -V | --version )    print_version ;;

            * )                 echo "Warning: $1 is not recognized."; Usage ;;
        esac
    fi
    shift
done

# Check the ulimit -n to set appropriate branch_count
ulimit_n=`ulimit -n`
tmp=$(( branch_count * 2 + 10 ))
if [ $ulimit_n -le $tmp ]; then
    tmp=$(( ulimit_n / 2 - 10 ))
    if [ $silent -eq 0 ]; then
        echo "Warning: -b $branch_count is not available on this system due to system limitation of simultaneously opened files $ulimit_n"
        echo "         -b $tmp will be used"
    fi
    branch_count=$tmp
fi

# Define a full path to cleanup utility
base_path=${BASH_SOURCE:-$0}
if [ "${base_path%%/*}" != "" ]; then 
    full_path=${PWD}"/"${base_path%/*}
else
    full_path=${base_path%/*}
fi


############################################
######             CLEANUP            ######
############################################
######           Total mode           ######
############################################
if [ $total_mode -eq 1 ]; then
    # Check user name
    if [ "$username" = "root" ]; then
        if [ $silent -eq 0 ]; then
            echo "Warning: -t total mode is not available for root"
        fi
        exit 1
    fi
    # Check host file
    if [ ! -n "$host_file" ]; then
        if [ $silent -eq 0 ]; then
            echo "Warning: Please set a host file to clean"
        fi
        exit 1
    fi
    if [ ! -f "$host_file" ]; then
        if [ $silent -eq 0 ]; then
            echo "Warning: $host_file doesn't exist"
        fi
        exit 1
    fi

    # Prework to define file count to split the host_file
    # Split up input file if necessary
    node_count=`cat $host_file | wc -l`
    quotient=$(( ( node_count + 1 ) / branch_count ))
    residue=$(( ( node_count + 1 ) % branch_count ))
    if [ $node_count -gt $branch_count ]; then
        line_count=$branch_count
        if [ $quotient -gt $branch_count ]; then
            file_count=$(( branch_count - 1 ))
        else
            file_count=$(( quotient - 1 ))
        fi

        csplit $host_file -s -f cleanup_input_${host_name}_${PPID}_ $line_count {$file_count}
        file_count=$(( file_count + 1 ))
    else
        file_count=1
    fi

    if [ $file_count -gt 1 ]; then
        # Hierarchical mode
        i=0
        while [ $i -le $file_count ]; do
            file_name=`printf "cleanup_input_${host_name}_${PPID}_%02d" $i`
            if [ $verbose -eq 1 ]; then
                echo "$file_name handling"
            fi
            # Running of cleanup
#            head=`cat $file_name | grep -v "^#" | head -n 1`
#            if [ $verbose -eq 1 ]; then
#                echo "    $rsh_cmd $head \"${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}\""
#            fi
#            if [ $parallel_mode -eq 1 ]; then
#                $rsh_cmd $head "${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}"&
#                i=$[i + 1]
#            else
#                $rsh_cmd $head "${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}"
#                i=$[i + 1]
#            fi
            while read line; do
                # Running of cleanup
                host=`echo $line | grep -v "^#"`
                if [ -z $host ]; then
                    continue
                fi
                if [ $verbose -eq 1 ]; then
                    echo "    $rsh_cmd $host \"${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}\""
                fi
                if [ $parallel_mode -eq 1 ]; then
                    $rsh_cmd $host "${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}"&
                    i=$(( i + 1 ))
                    break
                else
                    $rsh_cmd $host "${full_path}/${base_name} $opts_to_propagate -f ${PWD}/${file_name}"
                    if [ $? -eq 0 ]; then
                        i=$(( i + 1 ))
                        break
                    fi
                fi
            done < $file_name
        done
        if [ $parallel_mode -eq 1 ]; then
            wait
        fi
        echo $host_file | grep cleanup_input_ >/dev/null 2>&1
        if [ $? -eq 0 ]; then
            rm -rf $host_file
        fi
    else
        # Plain mode
        content=`cat $host_file | grep -v "^#"`
        for host in $content; do
            if [ $verbose -eq 1 ]; then
                echo "$rsh_cmd $host kill -9 -1"
            fi
            if [ $silent -eq 1 ]; then
                $rsh_cmd $host "kill -9 -1 >/dev/null 2>&1"&
            else
                $rsh_cmd $host "kill -9 -1"&
            fi
        done
        wait
        echo $host_file | grep cleanup_input_ >/dev/null 2>&1
        if [ $? -eq 0 ]; then
            rm -rf $host_file
        fi
    fi
    exit 0
fi


############################################
######          Common mode           ######
############################################
# Define an input_file previously generated by Hydra Process Mamager
if [ -z $username ] ; then
    username=`whoami`
fi
if [ -x $input_file ]; then
    if [ -n "$I_MPI_TMPDIR" ]; then
        input_file=${I_MPI_TMPDIR}/mpiexec_${username}_${PPID}.log
    else
        if [ -n "$TMPDIR" ] ; then
            input_file=${TMPDIR}/mpiexec_${username}_${PPID}.log
        else
            input_file=/tmp/mpiexec_${username}_${PPID}.log
        fi
    fi
fi
if [ ! -f "$input_file" ]; then
    if [ $silent -eq 0 ]; then
        echo "Warning: $input_file doesn't exist"
    fi
    exit 1
fi

# Prework to define file count to split the input_file
# Split up input file if necessary
node_count=`cat $input_file | wc -l`
quotient=$(( ( node_count + 1 ) / branch_count ))
residue=$(( ( node_count + 1 ) % branch_count ))
if [ $node_count -gt $branch_count ]; then
    line_count=$branch_count
    if [ $quotient -gt $branch_count ]; then
        file_count=$(( branch_count - 1 ))
    else
        file_count=$(( quotient - 1 ))
    fi

    csplit $input_file -s -f cleanup_input_${host_name}_${PPID}_ $line_count {$file_count}
    file_count=$(( file_count + 1 ))
else
    file_count=1
fi

if [ $file_count -gt 1 ]; then
    # Hierarchical mode
    i=0
    while [ $i -le $file_count ]; do
        file_name=`printf "cleanup_input_${host_name}_${PPID}_%02d" $i`
        if [ $verbose -eq 1 ]; then
            echo "$file_name handling"
        fi
        # Running of cleanup
#        head=`head -n 1 $file_name | sed -e "s/ .*//"`
#        if [ $verbose -eq 1 ]; then
#            echo "    $rsh_cmd $head \"${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}\""
#        fi
#        if [ $parallel_mode -eq 1 ]; then
#            $rsh_cmd $head "${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}"&
#            i=$[i + 1]
#        else
#            $rsh_cmd $head "${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}"
#            i=$[i + 1]
#        fi

        while read line; do
            # Running of cleanup
            host=`echo $line | sed -e "s/ .*//"`
            if [ $verbose -eq 1 ]; then
                echo "    $rsh_cmd $host \"${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}\""
            fi
            if [ $parallel_mode -eq 1 ]; then
                $rsh_cmd $host "${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}"&
                i=$(( i + 1 ))
                break
            else
                $rsh_cmd $host "ls >/dev/null" >/dev/null 2>&1
                if [ $? -ne 0 ]; then
                    continue
                fi
                $rsh_cmd $host "${full_path}/${base_name} $opts_to_propagate -i ${PWD}/${file_name}"&
                i=$(( i + 1 ))
                break
            fi
        done < $file_name
    done
    wait
    rm -rf $input_file
else
    # Plain mode
    while read line; do
        host=`echo $line | sed -e "s/ .*//"`
        pids=`echo $line | sed -e "s/${host} //"`
        if [ $verbose -eq 1 ]; then
            echo "    $rsh_cmd $host \"kill -9 $pids\""
        fi
        if [ $silent -eq 1 ]; then
            $rsh_cmd $host "kill -9 $pids >/dev/null 2>&1"&
        else
            $rsh_cmd $host "kill -9 $pids"&
        fi
        i=$(( i + 1 ))
    done < $input_file
    wait
    rm -rf $input_file
fi
exit 0
