set_nic_smp_affinity.sh

#!/bin/bash

#set -x

#for grep nic card pattern, you can change this one to meet your system environment!
NIC_PATTERN='em\|eth\|p1'
RX_PATTERN='rx-'

LS='/bin/ls'
CAT='/bin/cat'
GREP='/bin/grep'
PS='/bin/ps'
WC='/usr/bin/wc'

#TODO:add support to ubuntu system stop a service
#this time only support Redhat system 
SERVICE='/sbin/service'
CHKCONFIG='/sbin/chkconfig'

DISABLE_IRQBALANCE='NO'
ONLY_SHOW='NO'
QUIET='NO'
LOG_MSG='NO'

SYSNET_PATH='/sys/class/net'
QUEUE_PATH='queues'
RPS_ENTRIES_FILE="/proc/sys/net/core/rps_sock_flow_entries"
RPS_ENTRIES_NUM=32768
RPS_FLOW_CNT=4096

CPU_IDS=''
CPU_COUNT=0
ALL_USED_INTERFACE=''
SMP_AFFINITY_INF=''
RPS_RFS_INF=''
RPS_RFS_MASK=''


function parser_args() {
    local all_run=0
    while test -n "$1"; do
        case "$1" in
        -o|--only_show)
            DISABLE_IRQBALANCE='NO'
            ONLY_SHOW='YES'
            shift
            ;;
        -q|--quiet)
            QUIET='YES'
            shift
            ;;
        -a|--all)
            all_run=1
            shift
            ;;
        -b|--both)
            both=1
            shift
            ;;
        -d|--disable_irqbalance)
            DISABLE_IRQBALANCE='YES'
            shift
            ;;
        -n|--nic)
            ALL_USED_INTERFACE=$2
            shift 2
            ;;
        -h|--help)
            usage
            return 1
            ;;
        *)
            echo "Unknown argument: $1"
            usage
            return 1
            ;;
        esac
    done
    
    if [ $all_run -eq 0 ] && [ -z $ALL_USED_INTERFACE ]; then
        usage
        return 1
    fi  
    
    ###
    # suck !!!!
    if [ "$ONLY_SHOW" == 'YES' ] && [ "$QUIET" == 'YES' ]; then
        echo '$ONLY_SHOW:'$ONLY_SHOW
        echo '$QUIET:'$QUIET
        echo "[ERROR] You can set either -o or -q, but can not both !"
        return 1
    else
        return 0
    fi
    
    return 1
}

function debug_echo() {
    [ "$QUIET" == 'NO' ] && echo "$*"
    [ "$LOG_MSG" == 'YES' ] && logger "$*"
}

function version_compare () {
    if [[ $1 == $2 ]]
    then
        return 0
    fi
    local IFS=.
    local i ver1=($1) ver2=($2)
    # fill empty fields in ver1 with zeros
    for ((i=${#ver1[@]}; i<${#ver2[@]}; i++))
    do
        ver1[i]=0
    done
    for ((i=0; i<${#ver1[@]}; i++))
    do
        if [[ -z ${ver2[i]} ]]
        then
            # fill empty fields in ver2 with zeros
            ver2[i]=0
        fi
        if ((10#${ver1[i]} > 10#${ver2[i]}))
        then
            return 1
        fi
        if ((10#${ver1[i]} < 10#${ver2[i]}))
        then
            return 2
        fi
    done
    return 0
}


function check_interface_inuse() {
    if [ ${#@} -lt 1 ]; then
        return 1
    fi
    if [ -f "$SYSNET_PATH/$1/operstate" ]; then 
        local updown=$($CAT "$SYSNET_PATH/$1/operstate")
        if [ $? -eq 0 ]; then
            if [ "$updown" == "up" ];then
                return 0
            else
                return 1
            #return  "$updown" == "up"  ? 0 : 1
            fi
        fi
    fi
    return 1
}

function get_all_inused_interfaces() {
    local allInterFaces=$($LS "$SYSNET_PATH" | $GREP $NIC_PATTERN)
    #local len_t=${#allInterFaces}
    local len_t=$(echo $allInterFaces | $WC -w)
    if [ $len_t -lt 1 ];then
        return 1
    fi
    
    for inf_t in $allInterFaces; do
        #echo $inf_t
        check_interface_inuse "$inf_t"
        if [ $? -eq 0 ]; then
            #echo $inf_t
            ALL_USED_INTERFACE=$ALL_USED_INTERFACE" "$inf_t
        fi
    done
    #echo $ALL_USED_INTERFACE
    return 0
}

# get all nic card that support smp_affinity or rfs/rfs ,if no one ,return error
function get_smp_affinity_nics() {
    if [ -z $1 ];then
        return 1
    fi
    local nic_t_queue_count=$($LS "$SYSNET_PATH/$1/$QUEUE_PATH" | $GREP -c $RX_PATTERN)
    if [ $nic_t_queue_count -gt 1 ]; then
        SMP_AFFINITY_INF=$SMP_AFFINITY_INF" "$1
    else
        RPS_RFS_INF=$RPS_RFS_INF" "$1
    fi
    local count_t=$(echo $SMP_AFFINITY_INF$RPS_RFS_INF|$WC -w)
    
    if [ $count_t -gt 0 ]; then
        return 0
    fi
    return 1
}


# check for irq_balance running
# if $DISABLE_IRQBALANCE == 'YES' then disable irqbalance service
function check_irqbalance_on() {
    IRQBALANCE_ON=$($PS ax | $GREP -v "grep" | $GREP -q "irqbalance"; echo $?)

    if [ "$IRQBALANCE_ON" == "0" ] ; then
        echo " WARNING: irq_balance is running and will"
        echo "          likely override this script's affinitization."
        echo "          Please stop the irq_balance service and/or execute"
        echo "          'killall irqbalance'" 
        if [ "$ONLY_SHOW" == 'YES' ];then
            return 0
        fi
        if [ "$DISABLE_IRQBALANCE" == 'YES' ]; then
            $SERVICE irqbalance stop
            if [ $? -eq 0 ]; then 
            return 0
            fi
        fi
    else
        return 0
    fi
    return 1
}

function get_cpuinfo() {
    CPU_IDS=$($CAT "/proc/cpuinfo" | $GREP processor | $GREP -v grep | cut -d: -f 2)    
    CPU_COUNT=$(echo $CPU_IDS | $WC -w)
    if [ $CPU_COUNT -lt 1 ];then
        return 1
    fi
    return 0
}

#
function get_rpfs_affinity_mask() {
    local t=$(($CPU_COUNT/4))
    if [ $t -lt 1 ]; then
        return 1
    fi
    
    for i in $(seq $t); do
        RPS_RFS_MASK=$RPS_RFS_MASK'f'
    done
    #local t1=$(($CPU_COUNT-$t))
    #if [ $t1 -gt 0 ]; then
    #   for i in $(seq $t1); do
    #       RPS_RFS_MASK="0"$RPS_RFS_MASK
    #   done
    #fi
    
    local ok_t=$(echo $RPS_RFS_MASK | $GREP -q 'f'; echo $?)
    return $ok_t                
}

function do_rpfs_enable() {
    nic_dev=$1
    
    debug_echo "dealing with interface :"$nic_dev

    local rps_cpus_file_t="$SYSNET_PATH/$nic_dev/$QUEUE_PATH/rx-0/rps_cpus"
    debug_echo "dealing with rps_cpus_file :"$rps_cpus_file_t
    if [ -f $rps_cpus_file_t ]; then
        local rps_cpus=$($CAT $rps_cpus_file_t | tr -d '0')
        if [ -z "$rps_cpus" ]; then
            local rps_cpus=$($CAT $rps_cpus_file_t)
        fi
        #compare current rps_cus to $RPS_RFS_MASK
        if [ "$rps_cpus" != "$RPS_RFS_MASK" ]; then
            if [ "$ONLY_SHOW" == 'NO' ]; then
                echo $RPS_RFS_MASK > $rps_cpus_file_t
                local rps_cpus_tt=$($CAT $rps_cpus_file_t | tr -d '0')
                if [ "$rps_cpus_tt" != "$RPS_RFS_MASK" ]; then
                    echo "  [ERROR]:setting $rps_cpus_file_t to $RPS_RFS_MASK error !"
                    return 1
                fi
            fi
            debug_echo "    origin value :"$rps_cpus
            debug_echo "        will be set to :"$RPS_RFS_MASK
        else
            debug_echo "    already been set to:"$RPS_RFS_MASK
        fi
    else
        return 1
    fi
    
    local rps_flow_cnt_file_t="$SYSNET_PATH/$nic_dev/$QUEUE_PATH/rx-0/rps_flow_cnt"
    debug_echo "dealing with rps_flow_cnt_file :"$rps_flow_cnt_file_t
    if [ -f $rps_flow_cnt_file_t ]; then
        local rps_flow_cnt=$($CAT $rps_flow_cnt_file_t)
        #compare current rps_flow_cnt to $RPS_FLOW_CNT
        if [ "$rps_flow_cnt" != "$RPS_FLOW_CNT" ];then
            #   echo $RPS_FLOW_CNT 
            if [ "$ONLY_SHOW" == 'NO' ];then 
                echo $RPS_FLOW_CNT > $rps_flow_cnt_file_t
                local rps_flow_cnt_tt=$($CAT $rps_flow_cnt_file_t)
                if [ $rps_flow_cnt_tt != $RPS_FLOW_CNT ]; then
                    echo "  [ERROR]:setting $rps_flow_cnt_file_t to $RPS_FLOW_CNT error !"
                    return 1
                fi
            fi
            debug_echo "    origin value :" $rps_flow_cnt
            debug_echo "        will be set to :"$RPS_FLOW_CNT
        else
            debug_echo "        already been set to:"$RPS_FLOW_CNT
        fi
    else
        return 1
    fi
    return 0    
}

function do_rps_entries_enable() {
    debug_echo "dealing with rps_entries :"$RPS_ENTRIES_FILE
    local rps_entries=$($CAT $RPS_ENTRIES_FILE)
    if [ $rps_entries -eq $RPS_ENTRIES_NUM ]; then
        return 0
    else
        if [ "$ONLY_SHOW" == 'NO' ];then
            echo $RPS_ENTRIES_NUM > $RPS_ENTRIES_FILE
            local rps_entries_t=$($CAT $RPS_ENTRIES_FILE)
            if [ $rps_entries -eq $RPS_ENTRIES_NUM ]; then
                return 0
            else
                echo "  [ERROR]:setting $RPS_ENTRIES_FILE to $RPS_ENTRIES_NUM error !"
                return 1
            fi
        else
            debug_echo "    origin value :"$rps_entries
            debug_echo "    will be set to :"$RPS_ENTRIES_NUM
        fi
    fi
}

function do_smpaffinity_enable() {
    nic_dev=$1
    softIrqs=$($CAT "/proc/interrupts" | $GREP $nic_dev"-" | tr -s " " ":" | cut -d: -f 2)      
    softIrqCount=$(echo $softIrqs | $WC -w)
    
    #first=$(echo $softIrqs|cut -d" " -f 1) 
    if [ $CPU_COUNT -lt $softIrqCount ]; then
        debug_echo "Notice, CPU_COUNT:" $CPU_COUNT " is little than nic softIrqCount :" $softIrqCount
    fi
    
    debug_echo "$nic_dev nic softIrq list: "$softIrqs
    debug_echo "Total cpu list: "$CPU_IDS
    #debug_echo "First softIqr: "$first
    
    tmpC=$CPU_COUNT
    
    for s in $softIrqs; do
        tmpC=$(($tmpC - 1))
        if [ $tmpC -lt 0 ]; then
            tmpC=$(($CPU_COUNT-1))
        fi
        
        debug_echo "irq no. : "$s;
        local irq_t="/proc/irq/$s/smp_affinity_list"
        if [ -f $irq_t ]; then
            orig=$(cat $irq_t);
            debug_echo "    original set :"$orig 
            cpuid=$tmpC
            if [ "$orig" == "$cpuid" ]; then
                debug_echo "    already set to :"$cpuid
            else
                debug_echo "    will modified to: "$cpuid
                if [ "$ONLY_SHOW" == 'NO' ]; then
                    echo $cpuid > $irq_t;
                    cid=$(cat "$irq_t");
                    if [ "$cpuid" == "$cid" ];then
                        debug_echo "    done"
                    else
                        debug_echo "    Error setting,please check manually!!"
                        debug_echo "    check path:$irq_t"
                        return 1
                    fi
                fi
            fi
        fi
    done
    return 0
}

function usage() {
    echo "This Script will set smp_affinity or rps/rps automatically"
    echo "  via modify '/sys/class/net/' files and/or '/proc/' sysfs"
    echo "USAGE:$0 -a -o/-q"
    echo "  -a/--all:   Check all in_used interfaces"
    echo "  -b/--both:  Both set rps/rqs and smp irq affinity"
    echo "  -n/--nic:   Only check/set the specified nic interface"
    echo "  -d/--disable_irqbalance:    Disable irq_balance service automatically"
    echo "  -o/--only_show: Only show what will happen"
    echo "  -q/--quiet: Nothing will display if no error happen" 
    exit 1
}
 
############################MAIN################################

#if [ "$#" -lt 2 ]; then
#   usage
#   exit
#fi

#if [ "$1" == "-nic" ];then
#   nic_dev=$2
#else
#   usage
#fi

#echo "$*"
if [ ${#@} -lt 1 ]; then
    usage
fi

parser_args $*
if [ $? -gt 0 ]; then
    exit
fi


check_irqbalance_on
if [ $? -gt 0 ]; then
    exit
fi

if [ $(echo $ALL_USED_INTERFACE | $WC -w) -lt 1 ]; then
    get_all_inused_interfaces
    if [ $? -gt 0 ]; then
        exit
    fi
fi

get_cpuinfo
if [ $? -gt 0 ]; then
    exit
fi

debug_echo "All used interface :" $ALL_USED_INTERFACE
for inf in $ALL_USED_INTERFACE; do
    get_smp_affinity_nics $inf
done

debug_echo '    $SMP_AFFINITY_INF:'$SMP_AFFINITY_INF
debug_echo '    $RPS_RFS_INF:'$RPS_RFS_INF

if [ $(echo $SMP_AFFINITY_INF | $WC -w) -gt 0 ]; then
    for inf_t in $SMP_AFFINITY_INF; do
        do_smpaffinity_enable $inf_t
        if [ $? -gt 0 ]; then
            exit 1
        fi
    done
fi

if [ $(echo $RPS_RFS_INF | $WC -w) -gt 0 ]; then
    get_rpfs_affinity_mask
    if [ $? -gt 0 ]; then
        exit 1
    else
        do_rps_entries_enable
        for inf_t in $RPS_RFS_INF; do
            do_rpfs_enable $inf_t
            if [ $? -gt 0 ]; then
                exit 1
            fi
        done
    fi
fi
    
    
exit 1