#!/usr/bin/sh

# ve-set-hugepages
# Copyright (C) 2021 NEC Corporation
# This file is part of VE system initialization
#
# VE system initialization is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# VE system initialization is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public
# License along with VE system initialization; if not, see
# <http://www.gnu.org/licenses/>

MEMBIND=0
ADDITIONAL_HUGE_PAGES=0
# OVERCOMMIT_HUGE_PAGES is initialized as empty.
OVERCOMMIT_HUGE_PAGES=
SKIP_SETTING=0

SHOW=0
VARBOSE=0
MEMBIND_TMP=
ADDITIONAL_HUGE_PAGES_TMP=
OVERCOMMIT_HUGE_PAGES_TMP=

DEV_ID_VE1="1bcf:001c"
DEV_ID_VE3="1bcf:0039"
VE_ID="$DEV_ID_VE1
$DEV_ID_VE3"

CONF_DIR=/etc/opt/nec/ve/veos
CONF=/etc/opt/nec/ve/veos/ve-hugepages.conf

LOG1=:
LSPCI=/usr/sbin/lspci 

function help(){
        echo "Usage: $0 [-f filepath] [-a pages] [-o pages] [-m] [-s] [-vh]"
        echo -e "\t -a\t Add <pages> to the calculated Hugepages. "
        echo -e "\t -f\t Specify a path to a configuration file."
        echo -e "\t   \t When user specifies a file, the default configuration file is not loaded."
        echo -e "\t -m\t Specify a mode for NQSV socket scheduling with membind policy."
        echo -e "\t -o\t Set <pages> to the overcommit Hugepages."
        echo -e "\t -s\t Show the current values."
        echo -e "\t -v\t Verbose mode"
        echo -e "\t -h\t Show this help message."
        exit 0
}


# Calculate memory with PPS, MEMBIND condition
# The argument must be a number
function calc_pps_membind_cond(){
TMP_F=$1
if [ "$MEMBIND" = 1 ]; then
	let TMP_F=$TMP_F*$VH_NUMA_NUM
fi
if [ "$USE_PPS" = 1 ]; then
	let TMP_F=$TMP_F*2
fi
echo $TMP_F
}

#
# Calculate the memory size for PPS
# Return value: size in MB
function count_pps_mem(){

    local TOTAL_MEM=0

    # The routine is the same as the VEOS's PPS configuration loading 
    for i in $VE_SLOT; do
        local SLOT_N=${i#*veslot}
        local PPS_F_PATH=""
        local PPS_F_SZ=0
        local PPS_M_SZ=0
        local FILES="$CONF_DIR/ve-os-launcher.d/veswap.options $CONF_DIR/ve-os-launcher.d/$SLOT_N/veswap.options"
        local SEDCMD='/^ve-os-launcher@\*=/s/^ve-os-launcher@\*=\(.*\)/\1/gp'
        local OPTIONS=`sed -n -e $SEDCMD $FILES 2>/dev/null`
        for OPT in $OPTIONS; do
            case $OPT in
                --ve-swap-file-path=*)
                    PPS_F_PATH=${OPT##*=}
                    ;;
                --ve-swap-file-max=*)
                    PPS_F_SZ=${OPT##*=}
                    ;;
                --ve-swap-mem-max=*)
                    PPS_M_SZ=${OPT##*=}
                    ;;
            esac
        done

        if [ -n "$PPS_F_PATH" -a -n $"PPS_F_SZ" ]; then
            let TOTAL_MEM=$TOTAL_MEM+64
        fi
        if [ -n "$PPS_M_SZ" ]; then
            let TOTAL_MEM=$TOTAL_MEM+"$PPS_M_SZ"
        fi
    done
    echo $TOTAL_MEM             # Return value
}

#
# Set specified Hugepages
# $1   New value (pages)
# $3   Path to /proc/sys/vm/(nr_hugepages, nr_overcommit_hugepages)
# Return 0: Success
#        1: Success but the Hugepages is less than specified value
function set_pages(){
	CURRENT=`cat $2`
	if [ $1 != $CURRENT ]; then
		echo $1 > $2
		NEW_VAL=`cat $2`
		if [ $1 != $NEW_VAL ]; then
			return 1
		fi
	fi
	return 0
}

#
# Main part
#

if [ ! ${EUID:-${UID}} = 0 ]; then
       echo "$0: Permission denied"  1>&2
       exit 2
fi

while getopts a:f:o:msvh OPT
do
        case $OPT in
                "a" ) ADDITIONAL_HUGE_PAGES_TMP=$OPTARG ;;
                "f" ) CONF=$OPTARG ;;
                "m" ) MEMBIND_TMP=1 ;;
                "o" ) OVERCOMMIT_HUGE_PAGES_TMP=$OPTARG ;;
                "s" ) SHOW=1 ;;
                "v" ) VARBOSE=1 ;;
                "h" ) help ;;
        esac
done
shift $((OPTIND - 1))

source $CONF

if [ "$SKIP_SETTING" = "YES" ]; then
	exit 0
fi

CUR_HUGE_NUM=`cat /proc/sys/vm/nr_hugepages`
CUR_OVER_HUGE_NUM=`cat /proc/sys/vm/nr_overcommit_hugepages`

if [ $SHOW -eq 1 ]; then
	echo -e "nr_hugepages:$CUR_HUGE_NUM"
	echo -e "nr_overcommit_hugepages:$CUR_OVER_HUGE_NUM"
	exit 0
fi

if [ -n "$MEMBIND_TMP" ]; then 
	MEMBIND=$MEMBIND_TMP 
elif [  "$MEMBIND" = "YES" ]; then
	MEMBIND=1
fi
if [ -n "$ADDITIONAL_HUGE_PAGES_TMP" ]; then 
	ADDITIONAL_HUGE_PAGES=$ADDITIONAL_HUGE_PAGES_TMP
fi
if [ -n "$OVERCOMMIT_HUGE_PAGES_TMP" ]; then 
	OVERCOMMIT_HUGE_PAGES=$OVERCOMMIT_HUGE_PAGES_TMP
fi

if [ $VARBOSE -eq 1 ]; then
        LOG1=""
fi

#
# Environment check part
#
$LOG1 echo -e "[VE]"

# Test VE slots
VE_DEV=`$LSPCI -n | grep -F "$VE_ID"  | wc -l`
VE_VFIO_NUM=`$LSPCI -k -d $DEV_ID_VE3 | grep vfio-pci | wc -l`
VE_SLOT_NUM=`ls /dev/veslot* 2>/dev/null | wc -l`
TIMEOUT=60
while [ "$TIMEOUT" -ge 0 ]
do
	let VE_SUM=$VE_VFIO_NUM+$VE_SLOT_NUM
	if [ "$VE_SUM" != "$VE_DEV" ]; then
		# Wait /dev/veslot*
		sleep 1
		let TIMEOUT=$TIMEOUT-1
		VE_SLOT_NUM=`ls /dev/veslot* 2>/dev/null | wc -l`
	else
		break
	fi
done
if [ "$VE_VFIO_NUM" = "$VE_DEV" ]; then
	$LOG1 echo -e "All VE cards are bound to virtual machine."
	exit 0
fi
if [ "$VE_SLOT_NUM" = 0 ]; then
        echo "/dev/veslot<N> is not available. " 1>&2
	exit 2
fi
if [ "$TIMEOUT" -lt 0 ]; then
        echo "Some /dev/veslot are not available. Dev:$VE_DEV Slot:$VE_SLOT_NUM " 1>&2
        echo "Set Hugepages for the $VE_SLOT_NUM VEs" 1>&2
fi

# Count VE cards
VE_NUM=`ls /dev/veslot* -1 | wc -l`
$LOG1 echo -e "  VE cards\t\t  : $VE_NUM"

# Count VE cards per architecture
VE_NUM_VE1=`$LSPCI -d $DEV_ID_VE1 | wc -l`
let VE_NUM_VE3=$VE_NUM-$VE_NUM_VE1
$LOG1 echo -e "      VE1\t\t  : $VE_NUM_VE1" 
$LOG1 echo -e "      VE3\t\t  : $VE_NUM_VE3"

# Count VE cores
VE_TOTAL_CORE=0
VE_SLOT=`ls /dev/veslot*`
for i in $VE_SLOT
do
        VE_PATH=`/usr/bin/udevadm info -q path $i`
        let VE_TOTAL_CORE=$VE_TOTAL_CORE+`cat /sys/$VE_PATH/num_of_core`
done
$LOG1 echo -e "  Total VE cores\t  : $VE_TOTAL_CORE"


$LOG1 echo -e "[VH]"

# Get the number of VH NUMA nodes
VH_NUMA_NUM=`LANG=C /usr/bin/lscpu | grep NUMA.node.s |  awk '{print $3}'`
$LOG1 echo -e "  VH NUMA nodes\t\t  : $VH_NUMA_NUM"

# Get the number of VH cores
VH_CORE=`grep processor /proc/cpuinfo | wc -l`
$LOG1 echo -e "  VH Cores\t\t  : $VH_CORE"

# Get the total VH memory size (MB)
VH_MEM_TOTAL=`cat /proc/meminfo | grep MemTotal | awk '{$2=$2/1024;print int($2)}' `
$LOG1 echo -e "  VH mem(MB)\t\t  : $VH_MEM_TOTAL"

# Get the number of IB devices
IB_DEV=`$LSPCI | grep Infiniband | wc -l`
$LOG1 echo -e "  IB devices\t\t  : $IB_DEV"

# Get a virtualized environment 
VIRT_ENV=`systemd-detect-virt > /dev/null; echo $?`
if [ "$VIRT_ENV" -eq 0 ]; then
	$LOG1 echo -e "  Virtualized env.\t  : YES"
else
	$LOG1 echo -e "  Virtualized env.\t  : NO"
fi


$LOG1 echo -e "[Aurora software]"

# Get VEOS PPS configuration (MB)
USE_PPS=0
TOTAL_PPS_MEM=`count_pps_mem`
if [ "$TOTAL_PPS_MEM" -gt 0 ]; then
	USE_PPS=1
fi

# Accelerated IO (MB)
ACC_IO_MEM=8

# Check MPI library
USE_MPI=0
MPI_MEM=0
rpm -q nec-mpi-runtime > /dev/null
if [ $? -eq 0 ]; then
	USE_MPI=1
fi

# Check  ScaTeFS client
USE_SCATEFS=0
rpm -qa | grep scatefs-client > /dev/null
if [ $? -eq 0 ]; then
	USE_SCATEFS=1
fi
if [ "$USE_SCATEFS" = 1 -a "$VE_NUM" = 2 ]; then
	SCATEFS_MEM=64 #MB
elif [ "$USE_SCATEFS" = 1 -a "$VE_NUM" = 4 ]; then
	SCATEFS_MEM=128
else
	SCATEFS_MEM=0
fi


#
# Calculation part (MB)
#

TOTAL_MEM=0

#PPS
let TOTAL_MEM=$TOTAL_MEM+$TOTAL_PPS_MEM
$LOG1 echo -e "  Use VEOS PPS\t\t  : $USE_PPS"
$LOG1 echo -e "  Total PPS MEM(MB)\t  : $TOTAL_PPS_MEM"

#Accelerated IO
let TMP=$ACC_IO_MEM*$VE_TOTAL_CORE
TMP=`calc_pps_membind_cond $TMP`
let TOTAL_MEM=$TOTAL_MEM+$TMP
$LOG1 echo -e "  Total ACC-IO MEM(MB)\t  : $TMP"

# MPI library
TMP=0
if [ $USE_MPI = 1 ]; then
	let TMP=$VE_TOTAL_CORE*32
	if [ "$IB_DEV" -eq 0 ]; then
		let TMP=$TMP+$VH_CORE*64
	else
		let TMP=$TMP+$VH_CORE*16
	fi
	TMP=`calc_pps_membind_cond $TMP`
	let TOTAL_MEM=$TOTAL_MEM+$TMP
fi
$LOG1 echo -e "  Use NEC MPI\t\t  : $USE_MPI"
$LOG1 echo -e "  Total MPI MEM(MB)\t  : $TMP"

# ScaTeFS client
TMP=0
if [ $USE_SCATEFS = 1 ]; then
	let TMP=$VE_TOTAL_CORE*$SCATEFS_MEM
	if [ "$USE_PPS" = 1 ]; then
	        let TMP=$TMP*2
	fi
	let TOTAL_MEM=$TOTAL_MEM+$TMP
fi
$LOG1 echo -e "  Use ScaTeFS client\t  : $USE_SCATEFS"
$LOG1 echo -e "  Total ScaTeFS MEM(MB)   : $TMP"

$LOG1 echo -e "  Total memory for SW(MB) : $TOTAL_MEM"


#
# Setting part 
#

$LOG1 echo -e "OVERCOMMIT_HUGE_PAGES(-o) : $OVERCOMMIT_HUGE_PAGES"
$LOG1 echo -e "ADDITIONAL_HUGE_PAGES(-a) : $ADDITIONAL_HUGE_PAGES"

# Minimum pages
MIN=0
if [ ! "$VIRT_ENV" -eq 0 ]; then
	let MIN=513*$VE_NUM_VE3+256*$VE_NUM_VE1
fi
$LOG1 echo -e "Minimum pages\t\t  : $MIN"

# Maximum HugePages (For VM)
let MAX=$VH_MEM_TOTAL/4
if [ "$VIRT_ENV" -eq 0 ]; then
	$LOG1 echo -e "Maximum pages\t\t  : $MAX"
fi


RET_VAL=0

# Hugepages 
let TMP=$TOTAL_MEM/2+$ADDITIONAL_HUGE_PAGES
if [ $TMP -lt $MIN ]; then
	TMP=$MIN
fi
if [ "$VIRT_ENV" -eq 0 -a $TMP -gt $MAX ]; then
	let VH_MEM_TOTAL_2MB=$VH_MEM_TOTAL/2
	echo "The required Hugepages is not allocated, because the main memory" \
		"is too small." \
		"Required: $TMP, Total: $VH_MEM_TOTAL_2MB" 1>&2
	exit 2
fi

$LOG1 echo -e "Hugepages\t\t  : $TMP"
set_pages $TMP /proc/sys/vm/nr_hugepages
if [ $? -eq 1 ]; then
	echo "The current Hugepages is different from the required value. " \
		"Required: $TMP, Current: `cat /proc/sys/vm/nr_hugepages`" 1>&2
	RET_VAL=1
fi

# Overcommit
if [ -z "$OVERCOMMIT_HUGE_PAGES" ]; then
	# The half of VH memory
	let TOTAL_OVERCOMMIT=$VH_MEM_TOTAL/4 
else
	let TOTAL_OVERCOMMIT=$OVERCOMMIT_HUGE_PAGES
fi
set_pages $TOTAL_OVERCOMMIT /proc/sys/vm/nr_overcommit_hugepages
if [ $? -eq 1 ]; then
	echo "The current overcommit Hugepages is different from the required value." \
		"Required: $TOTAL_OVERCOMMIT, " \
		"Current: `cat /proc/sys/vm/nr_overcommit_hugepages`" 1>&2
	RET_VAL=1
fi
$LOG1 echo -e "Overcommit Hugepages\t  : $TOTAL_OVERCOMMIT"


exit $RET_VAL


