#!/bin/bash
# Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
# Description: dataturbo
# Create: 2022/01/22
#

CUR_DIR="$( cd "$( dirname "$0"  )" && pwd )"
readonly DATATURBO_SERVICE="dataturbo.service"
DATATURBO_VERSION="1.0.0.SPC18"
DATATURBO_PATH=/opt/oceanstor/dataturbo
DATATURBO_CONF_PATH=/opt/oceanstor/dataturbo/conf
RAND_NUM=${RANDOM}
TEMP_BAK_DIR=/tmp/oceanstor_dataturbo_bak_${RAND_NUM}_$$
UPGRADE_RESULT=${DATATURBO_PATH}/script/upgrade_res
readonly DPC_USER_CONF="${DATATURBO_PATH}/conf/dpc_user_config.xml"
readonly SHIM_CONF="${DATATURBO_PATH}/conf/shim_config.xml"
readonly LOCK_FILE="/var/log/upgrade.shell.lock"

CGROUP_LIMIT=( 
"8,30" 
"12,50"
"20,95"
)

# 顺序停服务,后续扩展直接往列表后面添加服务即可。
readonly stop_srv_array=(dataturbo.service)

readonly LOG_FILE="/var/log/dataturbo/dataturbo_script/dataturbo.action.log"

function usage()
{
    echo "usage: $0 [-h|-f]"
    echo "  -h                  display help."
    echo "  -f                  force upgrade dataturbo client (Force upgrade will stop the service, and the upgrade rollback may fail)."
}

function log_info()
{
    if  [ ! -f ${LOG_FILE} ] || [ -L ${LOG_FILE} ]; then
        return
    fi
    local func_name=${FUNCNAME[1]}
    if [ "$2" != "" ]; then
        func_name="$2"
    fi
    echo "[`date "+%Y-%m-%d %T"`:$$][INFO][dataturbo][${0##*/}][${func_name}]:" "$1" >> ${LOG_FILE}
}

function log_error()
{
    if  [ ! -f ${LOG_FILE} ] || [ -L ${LOG_FILE} ]; then
        return
    fi
    local func_name=${FUNCNAME[1]}
    if [ "$2" != "" ]; then
        func_name="$2"
    fi
    echo "[`date "+%Y-%m-%d %T"`:$$][ERR][dataturbo][${0##*/}][${func_name}]:" "$1" >> ${LOG_FILE}
}

function log_warn()
{
    if  [ ! -f ${LOG_FILE} ] || [ -L ${LOG_FILE} ]; then
        return
    fi

    local func_name=${FUNCNAME[1]}
    if [ "$2" != "" ]; then
        func_name="$2"
    fi
    echo "[`date "+%Y-%m-%d %T"`:$$][WARN][dataturbo][${0##*/}][${func_name}]:" "$1" >> ${LOG_FILE}
}

function log_info_print()
{
    echo -e "$1"
    log_info "$1" "${FUNCNAME[1]}"
}

function log_error_print()
{
    echo -e -n "\033[31m[ERROR]\033[0m"; echo -e "$1"
    log_error "$1" "${FUNCNAME[1]}"
}

function log_warn_print()
{
    echo -e -n "\033[31m[WARN]\033[0m"; echo -e "$1"
    log_warn "$1" "${FUNCNAME[1]}"
}

#磁盘空间检查函数
function do_check_free_space()
{
    SPACE_STR=$(df -k "$1" | sed -n '$p')
    if [ "$SPACE_STR" = "" ]; then
        log_error_print "Check free space failed"
        return 1
    fi

    FREE_SPACE=$(df -k "$1" | sed -n '$p' | awk '{print $(NF-2)}')
    if [ -z "$FREE_SPACE" ] || [ "$FREE_SPACE" -lt "$2" ]; then
        log_error_print "[$1] File Directory need more disk free space, need[$(expr $2 / 1024)M], The current disk space is [$(expr ${FREE_SPACE} / 1024)M]"
        return 1
    fi
}

# 内存空间检查  2G 210241024 -->2097152
function do_check_mem_space()
{
    local spaceStr=""
    if [ -z "$1" ]; then
        log_error_print "Param error[$1]"
        return 1
    fi

    spaceStr=$(cat /proc/meminfo |awk '/MemAvailable/{print $2}')
    if [ "${spaceStr}" = "" ]; then
        log_error_print "Get memory free space failed"
        return 1
    fi

    if [ "${spaceStr}" -lt "$1" ]; then
        let freeSpaceMB=$(expr "$1" / 1024 )
        log_error_print "More memory free space need["${freeSpaceMB}"MB],The current memory size is [$(expr "${spaceStr}" / 1024)MB]"
        return 1
    fi
}

function check_old_rpm_package()
{
    package_name=$(ls ${DATATURBO_PATH}/packages/oceanstor_dataturbo*.rpm 2>/dev/null)
    if [ ! ${package_name} ]; then
        log_warn_print "Can not find oceanstor dataturbo rpm backup package."
        return 1
    fi
    return 0
}

function backup_file()
{
    mkdir -m 700 ${TEMP_BAK_DIR}
    if [ $? -ne 0 ]; then
        log_error_print "Create backup dir failed."
        return 1
    fi

    #backup
    cp -rp ${DATATURBO_PATH}/packages ${TEMP_BAK_DIR}/
    cp -rp ${DATATURBO_PATH}/conf ${TEMP_BAK_DIR}/
    cp -p ${DATATURBO_PATH}/bin/dpc_username_config.ini ${TEMP_BAK_DIR}/
    return 0;
}

function upgrade_rollback()
{
    old_package=$(ls ${TEMP_BAK_DIR}/packages/)
    if [ "${old_package}" == "" ]; then
        log_warn_print "Cat not find the old package"
        log_error_print "Rollback failed."
        return 1
    fi
    log_warn_print "Upgrade failed, start rollback to ${old_package}"
    # 卸载环境上的dataturbo包
    cur_env_packages=$(rpm -qa|grep dataturbo-*)
    for package in ${cur_env_packages}; do
        rpm -e ${package} --noscripts --nodeps 2>/dev/null
    done
    rm -rf ${DATATURBO_PATH}
    
    # 安装rpm老包
    user_name=$(cat ${TEMP_BAK_DIR}/dpc_username_config.ini | grep "username" | cut  -f2  -d "=")
    CUSTOM_USER=${user_name} rpm -ivh --nodeps ${TEMP_BAK_DIR}/packages/oceanstor_dataturbo*.rpm
    if [ $? -ne 0 ]; then
        log_error_print "Install dataturbo rpm failed."
        log_error_print "Rollback failed."
        return 1
    fi
    cp -af ${TEMP_BAK_DIR}/packages/oceanstor_dataturbo*.rpm ${DATATURBO_PATH}/packages/

    # 恢复配置文件
    rm -rf ${DATATURBO_PATH}/conf
    cp -rp ${TEMP_BAK_DIR}/conf ${DATATURBO_PATH}/
    rm -f ${DATATURBO_PATH}/bin/dpc_username_config.ini
    cp -p ${TEMP_BAK_DIR}/dpc_username_config.ini ${DATATURBO_PATH}/bin/

    log_info_print "Rollback complete."
    return 0
}

function set_cpu_limit()
{
    if [ ! -f "${DPC_USER_CONF}" ]; then
        echo "ERROR: ${DPC_USER_CONF} is not exsit, please reinstall."
        exit 1
    fi

    sed -i "s/<DpcCgroupLimitCPU min=\"1\" max=\"95\">[0-9]*<\/DpcCgroupLimitCPU>/<DpcCgroupLimitCPU min=\"1\" max=\"95\">$1<\/DpcCgroupLimitCPU>/g" "${DPC_USER_CONF}" >/dev/null 2>&1
}

function set_mem_limit()
{
    if [ ! -f "${DPC_USER_CONF}" ]; then
        echo "ERROR: ${DPC_USER_CONF} is not exsit, please reinstall."
        exit 1
    fi

    sed -i "s/<DpcCgroupLimitMemory min=\"4\" max=\"20\">[0-9]*<\/DpcCgroupLimitMemory>/<DpcCgroupLimitMemory min=\"4\" max=\"20\">$1<\/DpcCgroupLimitMemory>/g" "${DPC_USER_CONF}" >/dev/null 2>&1
}

function set_cgroup_limit_conf()
{
    local str=${CGROUP_LIMIT[$1-1]}
    local memory=$(echo "${str}" | awk -F "," '{print $1}')
    local cpu=$(echo "${str}" | awk -F "," '{print $2}')

    set_cpu_limit ${cpu}
    set_mem_limit ${memory}
    log_info "set cpu[${cpu}%] and memory[${memory}GB] limit OK."
}

function upgrade_rpm()
{

    # 升级rpm包
    rpm -Uvh --nodeps ${CUR_DIR}/packages/oceanstor_dataturbo*.rpm
    if [ $? -ne 0 ]; then
        upgrade_rollback
        return 1
    fi

    upgrade_ret=$(cat ${UPGRADE_RESULT} 2>/dev/null)
    if [ "${upgrade_ret}" != "" ] && [ "${upgrade_ret}" == "1" ]; then
        upgrade_rollback
        return 1
    fi

    if [ ! -f "${TEMP_BAK_DIR}/conf/dpc_user_config.xml" ]; then
        dpc_level_val=$(cat ${SHIM_CONF} | grep DpcLevelVal)
        dpc_level_val=${dpc_level_val#*>}
        dpc_level_val=${dpc_level_val%%<*}
        set_cgroup_limit_conf ${dpc_level_val}
    fi

    # 更新 rpm 备份包
    rm -f ${DATATURBO_PATH}/packages/*
    cp -af ${CUR_DIR}/packages/oceanstor_dataturbo*.rpm ${DATATURBO_PATH}/packages/
    log_info_print "Upgrade dataturbo successfully."
    return 0
}

function stop_srv_array()
{
    local result=0

    for srv in "${stop_srv_array[@]}";
    do
        systemctl stop "${srv}" > /dev/null 2>&1
        if [ $? -ne 0 ]; then
            result=1
            log_error_print "Stop ${srv} failed. See \"systemctl status ${srv}\" for details."
        fi

        if [ "${result}" != "1" ]; then
            log_info_print "Stop ${srv} service successfully."
        fi
    done
    return $result
}

function stop_service()
{
    local result=0

    systemctl stop "${DATATURBO_SERVICE}" > /dev/null 2>&1
    if [ $? -ne 0 ]; then
        result=1
        log_error_print "Stop ${DATATURBO_SERVICE} failed. See \"systemctl status ${DATATURBO_SERVICE}\" for details."
    fi

    if [ "${result}" != "1" ]; then
        log_info_print "Stop dataturbo's service successfully."
        return ${result}
    fi
    return ${result}
}

function check_dpc_status()
{
    RETRY_SLEEP_TIME=0.5
    DPC_PROC="/opt/oceanstor/dataturbo/bin/dpc"
    # When check dpc status by ps cmd failed, it will be retried 3 times
    for i in {0..3}
    do
        ps -ef | grep ${DPC_PROC} | grep -v grep 1>/dev/null
        if [ $? -eq 0 ];then  # dpc is running
            log_info_print "dpc process is running."
            return 1
        fi
        sleep $RETRY_SLEEP_TIME
    done

    return 0
}

function stop_dpc_process()
{
    for i in {0..3} # When getting PID fails, it will be retried 3 times
    do
        DPC_PID=$(pidof "dpc")
        if [ "${DPC_PID}" ];then
            kill -9 ${DPC_PID} 2>/dev/null
            return 0
        fi
        sleep 0.5
    done
    log_warn "Get pid of dpc process failed, process not exist."
    return 1
}

function check_srv_status()
{
    # Declare and assign separately to avoid masking return values.
    local status_in_system
    local loaded
    local active

    status_in_system=$(systemctl status dataturbo.service 2>/dev/null)
    loaded=$(echo "${status_in_system}"| grep -c "Loaded: loaded")
    active=$(echo "${status_in_system}"| grep -c "Active: active")
    if [ "$loaded" -eq "0" ] || [ "$active" -eq "0" ]; then
        return 0
    else
        log_warn_print "dataturbo.service status is active"
        return 1
    fi
}

# 大于
function version_gt()
{
    test "$(echo "$@" | tr " " "\n" | sort -V | head -n 1)" != "$1";
}

#小于
function version_lt() 
{ 
    test "$(echo "$@" | tr " " "\n" | sort -rV | head -n 1)" != "$1"; 
}

# 大于等于
function version_ge()
{
    test "$(echo "$@" | tr " " "\n" | sort -rV | head -n 1)" == "$1";
}

# 升级前检查版本号
function check_dataturbo_version()
{
    cur_rpm_name=$(rpm -qa dataturbo 2>/dev/null)
    if [ "${cur_rpm_name}" == "" ]; then
        log_warn_print "DataTurbo has not been installed yet."
        return 1
    fi

    rpm_path=${CUR_DIR}/packages/
    dataturbo_rpm=$(ls $rpm_path | grep "dataturbo")
    DATATURBO_VERSION=${dataturbo_rpm%%-*}
    DATATURBO_VERSION=${DATATURBO_VERSION##*_}
    release=${dataturbo_rpm#*-}
    release=${release%%.*}
    cur_rpm_version=${cur_rpm_name#*-}
    cur_rpm_version=${cur_rpm_version%%-*}
    cur_rpm_release=${cur_rpm_name##*-}
    cur_rpm_release=${cur_rpm_release%%.*}
    log_info_print "Current package version: ${cur_rpm_version}-${cur_rpm_release}, upgrade package version: ${DATATURBO_VERSION}-${release}"
    # 当前版本小于升级版本直接升级
    if  version_lt ${cur_rpm_version} ${DATATURBO_VERSION}; then
        return 0
    # 如果当前版本大于升级版本返回1 
    elif version_gt ${cur_rpm_version} ${DATATURBO_VERSION}; then
        log_error_print "Current version is greater than or equal to upgrade version."
        return 1
    # 如果当前版本等于升级版本则比较release号(时间戳)
    else 
        if version_ge ${cur_rpm_release} ${release};then
            log_error_print "Current release version ${cur_rpm_release} is greater than or equal to upgrade release version ${release}."
            return 1
        fi
    fi
    return 0
}

# 检查用户配置文件是否存在，存在 升级后才可继承配置
function check_config_file()
{
    conf_list=(${DATATURBO_CONF_PATH}/cache_product_config_release.xml
               ${DATATURBO_CONF_PATH}/cache_product_config.xml
               ${DATATURBO_CONF_PATH}/command_download.xml
               ${DATATURBO_CONF_PATH}/dpc.conf
               ${DATATURBO_CONF_PATH}/client.cfg
               ${DATATURBO_CONF_PATH}/shim_config.xml
               ${DATATURBO_CONF_PATH}/upgrade_info.conf
               ${DATATURBO_CONF_PATH}/dpc_config.xml
               ${DATATURBO_CONF_PATH}/config/node_config.xml
               ${DATATURBO_CONF_PATH}/../bin/dpc_username_config.ini)

    cont=0   
    for file in ${conf_list[*]}; do
        ls $file 1>/dev/null 2>/dev/null
        if [ $? -ne 0 ]; then
            log_warn_print "${file} not exist"
            let cont++;
        fi
    done

    if [ ${cont} -ne 0 ]; then
        return 1;
    fi

    return 0;
}

# 升级前检查。
function upgrade_check()
{
    # 磁盘空间检查函数 500M
    do_check_free_space /opt 512000
    if [ "$?" = "1" ]; then
        log_error_print "Check free space failed."
        return 1
    fi

    # 内存空间检查 2G 2*1024*1024 -->2097152
    do_check_mem_space 2097152
    if [ "$?" = "1" ]; then
        log_error_print "Check mem space failed."
        return 1
    fi

    # 升级前检查版本号
    check_dataturbo_version
    if [ "$?" == "1" ]; then
        log_error_print "Check dataturbo version failed."
        return 1
    fi

    if [ "$1" != "-f" ]; then
        check_old_rpm_package
        if [ "$?" -ne 0 ]; then
            return 1
        fi

        check_config_file
        if [ "$?" -ne 0 ]; then
            log_warn_print "Check config file failed."
            return 1
        fi
    fi

    # 升级场景检查dataturbo服务状态.
    check_srv_status
    ret=$?
    if [ "$ret" == "1" ] && [ "$1" == "-f" ]; then
        # 停止服务
        stop_srv_array
        if [ "$?" == "1" ]; then
            log_error_print "Stop dataturbo.service failed."
            return 1
        fi
    elif [ "$ret" == "1" ] && [ "$1" != "-f" ]; then
        log_warn_print "Check services status error, dataturbo service should be stopped before upgrade."
        return 1
    fi

    # 检查dpc进程
    check_dpc_status
    ret=$?
    if [ "$ret" == "1" ] && [ "$1" == "-f" ]; then
        # 停止服务
        stop_dpc_process
        if [ "$ret" == "1" ]; then
            log_error_print "dpc process stop failed."
            return 1
        fi
    elif [ "$ret" == "1" ] && [ "$1" != "-f" ]; then
        log_warn_print "check process status error, dpc process should be stopped before upgrade."
        return 1
    fi

    return 0
}

function lock_shell()
{
    local thepid=""
    local ret=""

    if [ "$1" = "" ]; then
        log_warn_print "need lock file param"
        exit 1
    fi

    if [ -f "${1}" ]; then
        thepid=$(cat "${1}")
        if [ "${thepid}" != "" ]; then
            ret=$(ps ax |awk '{print $1}'| grep -c "^[[:blank:]]*${thepid}[[:blank:]]*$")
            if [ ${ret} -eq 1 ]; then
                log_warn_print "This shell task is already running, exit."
                exit 1
            fi
        fi
    fi

    echo $$ > "${1}"
    chmod 600 "${1}"
    return 0
}

function unlock_shell()
{
    if [ "$1" = "" ]; then
        return 1
    fi

    if [ -f "${1}" ]; then
        rm -f "${1}"
    fi

    return 0
}

function abnormal_exit()
{
    unlock_shell "${lockFile}"
    exit 1
}

function do_upgrade() {
    lock_shell "${LOCK_FILE}"
    log_info "do upgrade with option: [$1]"
    # 升级前检查,参数-f:会自动停止DataTurbo服务
    upgrade_check $1
    if [ $? -ne 0 ]; then
        log_error_print "Upgrade check failed."
        abnormal_exit
    fi

    backup_file
    if [ $? -ne 0 ]; then
        log_error_print "Backup file before upgrade failed."
        abnormal_exit
    fi

    upgrade_rpm
    ret=$?

    rm -rf ${TEMP_BAK_DIR}
    unlock_shell "${lockFile}"
    
    exit $ret
}

if [ $# -gt 1 ]; then
    log_error_print "Wrong input parameter."
    usage
    exit 1
fi

PARA="$1"
if [ -z "${PARA}" ]; then
    do_upgrade
fi

while getopts hf options 2>/dev/null
do
    case $options in
        h)  #显示帮助信息
            usage
            exit 0
        ;;
        f)
            do_upgrade "-f"
        ;;
        *)
            usage
            exit 1
            ;;
    esac
done
usage
exit 1