#!/usr/bin/bash
# Copyright (c) 2020 International Business Machines
#
# This program 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.
#
# This program 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 this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

# Simple script designed to save and restore the SMT state.

readonly SMT_STATE=/var/lib/powerpc-utils/smt.state
readonly SMT_STATE_BKP=/tmp/smt.state.backup
readonly PPC64CPU=/usr/sbin/ppc64_cpu

update_state() {
    # sanity check
    smt_state_exist

    # users can online and offline individual CPUs, creating an inconsistent
    # SMT state. In this case ppc64_cpu --smt displays a more complex
    # configuration, such case is beyond the scope here so we simply ignore it
    if [ $($PPC64CPU --smt -n | cut -d= -f2 | grep ^[1-8]$ > /dev/null; echo $?) -ne 0 ]
    then
        return
    fi

    # create a temporary backup
    cp -p $SMT_STATE $SMT_STATE_BKP
    if [ $? -ne 0 ]
    then
        logger -t ${0##*/}[$$] "Error saving $SMT_STATE_BKP"
        exit 1
    fi

    # update SMT state file
    local current_smt=$($PPC64CPU --smt -n | cut -d= -f2)
    sed -i "s/SMT_VALUE=[0-9]/SMT_VALUE=$current_smt/" $SMT_STATE
    if [ $? -ne 0 ]
    then
        logger -t ${0##*/}[$$] "Error updating $SMT_STATE"
        cp -p $SMT_STATE_BKP $SMT_STATE
        rm $SMT_STATE_BKP
        exit 1
    fi

    rm $SMT_STATE_BKP
}

set_smt() {
    # sanity check
    smt_state_exist

    # get saved SMT value and set it
    local smt=$(grep "^SMT_VALUE=[1-8]$" $SMT_STATE | cut -d= -f2)
    if [ "$smt" -lt 1 ] || [ "$smt" -gt 8 ]
    then
        logger -t ${0##*/}[$$] "$smt is an invalid SMT state"
        exit 1
    fi

    $PPC64CPU --smt="$smt"
    if [ $? -ne 0 ]
    then
        logger -t ${0##*/}[$$] "Error updating SMT=$smt"
        exit 1
    fi
}

smt_state_exist() {
    if [ ! -f $SMT_STATE ]
    then
        logger -t ${0##*/}[$$] "$SMT_STATE not found"
        exit 1
    fi

    grep "^SMT_VALUE=[1-8]$" $SMT_STATE > /dev/null 2>&1
    if [ $? -ne 0 ]
    then
        logger -t ${0##*/}[$$] "$SMT_STATE does not have SMT_VALUE defined"
        exit 1
    fi
}

check_smt_capable() {
    local msg=$($PPC64CPU --smt 2>&1 | grep -q "Machine is not SMT capable"; echo $?)
    if [ "$msg" -eq 0 ]
    then
        logger -t ${0##*/}[$$] "Machine is not SMT capable, exiting"
        exit 0
    fi
}

usage() {
    printf "$0 is a script designed to save/restore the SMT state.\n"
    printf "Usage: %s [--save | --load | --help]\n" "$0"
    printf "\t--save (default): save current SMT state\n"
    printf "\t--load          : set SMT to the value stored\n"
    printf "\t--help          : print this message\n"
    printf "SMT state is saved in $SMT_STATE\n"
}

main() {
    # if machine is not SMT capable don't even bother to continue
    check_smt_capable

    local action="$1"

    case "$action" in
        --load)
            set_smt
            ;;

        --save)
            update_state
            ;;

        --help)
            usage
            ;;

        *)
            logger -s -t ${0##*/}[$$] "Invalid action $1"
            usage
            exit 1
            ;;
    esac
}

action=${1:---save}
main "$action"
