#!/bin/bash

# Author: Robert E. Gilligan <gilligan@vyatta.com>
# Date: 2008
# Description: CLI back-end script to manipulate NIC interrupt CPU affinity.

# **** License ****
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
# 
# 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.
# 
# This code was originally developed by Vyatta, Inc.
# Portions created by Vyatta are Copyright (C) 2006, 2007, 2008 Vyatta, Inc.
# All Rights Reserved.
# **** End License ****

# Provides sub-commands to:
#   - Check the validity of an affinity mask value
#   - Set the affinity mask to the IRQs being used by an interface
#   - Reset the affinity mask of the IRQs being used by an interface to the
#     system default value of all-ones.
#   - Print the affinity mask of the IRQs being used by an interface
#
# If the NIC in question supports multiple IRQs, the "set" sub-command
# sets all IRQs to the same mask.  The "print" sub-command displays
# the mask of each IRQ individually.
#

# Max number of hex characters in an IRQ affinity mask.  Support up to 64 CPUs.
MAX_MASK=16

# Set up some global values...
numcpus=`grep -c -e "^processor" /proc/cpuinfo`
declare -i maxmask=(2**numcpus)
let maxmask=maxmask-1
maxmaskhex=`printf "%x" ${maxmask}`

print_usage()
{
    echo "Usage:"
    echo -e "\t$0 check <ifname> <mask>"
    echo -e "\t$0 set <ifname> <mask>"
    echo -e "\t$0 reset <ifname>"
    echo -e "\t$0 print <ifname>"
}

get_irqnums()
{
    irqnums=`grep $1 /proc/interrupts | awk -F ': ' '{ print $1 }'`
    if [ -z "$irqnums" ]; then
	echo "Invalid interface name: $1"
	return 1
    fi
    return 0
}


get_mask()
{
    mask=$1
    # mask must be a short hex value
    if [ ${#mask} -gt $MAX_MASK ]; then
	echo "mask too long: ${#2} characters."
	return 1
    fi

    # strip out all the hex digits
    exmask=`echo $mask | sed -e s/[0-9a-fA-F]//g`

    # if anything is left, its not hex
    if [ ! -z "$exmask" ]; then
	echo "Invalid characters in hex mask: $exmask"
	return 1
    fi

    declare -i intmask=0x${mask}

    # Make sure that mask holds at least one bit, and holds no more bits
    # than we have CPUs.

    if [ ${intmask} -eq 0 ]; then
	echo "Mask can not be 0."
	return 1
    fi

    if [ $intmask -gt $maxmask ]; then
	echo "Mask is too large.  Maximum hexidecimal bitmask is: ${maxmaskhex}"
	return 1
    fi

    return 0
}


#
# Don't waste our time with uniprocessor machines
#
check_uniproc()
{
    if [ $maxmask -eq 1 ]; then
	echo "This machine has only 1 CPU."
	echo "Can only set SMP affinity on multi-processor machines"
	return 1;
    fi
    return 0
}
	      

case "$1" in
    check)
	# Note: We don't validate the interface name even though
	# it is available as a command argument.  That is because 
	# the interface may not exist or may not be configured at
	# the time the check is performed.  
	#
	if [ $# -ne 3 ]; then
		print_usage
		exit 1
	fi

	if ! check_uniproc ; then
		exit 1
	fi

	if ! get_mask $3 ; then
	    exit 1
	fi
	exit 0
	;;

    set)
	if [ $# -ne 3 ]; then
		print_usage
		exit 1
	fi

	if ! check_uniproc ; then
		exit 1
	fi

	if ! get_irqnums $2 ; then
	    exit 1
	fi

	if ! get_mask $3 ; then
	    exit 1
	fi

	for irqnum in $irqnums ; do
	    echo $mask > /proc/irq/$irqnum/smp_affinity
	done

	if [ $? -ne 0 ]; then
	    echo "Couldn't assign smp_affinity.  Exit status: $?"
	    exit 1
	fi
	;;

    reset)
	if [ $# -ne 2 ]; then
		print_usage
		exit 1
	fi
	if ! get_irqnums $2 ; then
	    exit 1
	fi

	if [ -e /proc/irq/default_smp_affinity ]; then
	    defmask=`cat /proc/irq/default_smp_affinity`
	else
	    defmask=$maxmaskhex
	fi

	for irqnum in $irqnums ; do
	    echo $defmask > /proc/irq/$irqnum/smp_affinity
	    if [ $? -ne 0 ]; then
		echo "Couldn't assign smp_affinity for IRQ $irqnum.  Exit status: $?"
		exit 1
	    fi
	done
	;;
     

    print)
	if [ $# -ne 2 ]; then
		print_usage
		exit 1
	fi
	if ! get_irqnums $2 ; then
	    exit 1
	fi

	for irqnum in $irqnums ; do
	    mask=`cat /proc/irq/$irqnum/smp_affinity`
	
	    if [ -z $mask ]; then
		echo "Couldn't get smp_affinity for interface $2, irq $irqnum"
		exit 1
	    fi

	    echo "Interface: $2 IRQ: $irqnum Mask: $mask"
	done
	;;

    *)
	print_usage
	exit 1
	;;
	
esac