#!/bin/bash

set -euo pipefail

DRYRUN=0
INTERRUPTS=/proc/interrupts

# Local development/debug/testing:
# put files:
# - cat /proc/interrupts -> ./interrupts
# - lscpu --extended=CPU,SOCKET -> ./lscpu_output
# rss-ladder --test ixgbe eth1 1

if [ "${1:-}" = '--test' ]; then
	INTERRUPTS=interrupts
	lscpu() {
		cat lscpu_output
	}
	__driver__="$2"
	ethtool() {
		echo driver "$__driver__"
	}
	shift 2
	DRYRUN=1
fi


usage() {
	echo "Usage: $0 <dev> [cpu socket]"
	echo "Example: $0 eth1"
	echo "Example: $0 eth1 0"
	echo "Example: $0 eth2 1"
}

__socket_cpu_list() {
	local socket="$1"
	lscpu --extended=CPU,SOCKET | grep -v CPU | egrep "[0-9]+[ ]+$socket$" | awk '{print $1}'
}

cpus_list_all() {
	lscpu --extended=CPU,SOCKET | grep -v CPU | awk '{print $1}'
}

socket_cpu_list() {
	local socket="$1"
	local queue_count="$2"
	local cpus times=4
	cpus=( $(__socket_cpu_list $socket) )
	for _ in $(seq $times); do
		__socket_cpu_list $socket
	done > /tmp/cpus.$socket
	head -n $queue_count /tmp/cpus.$socket
	rm -f /tmp/cpus.$socket
}

interrtupts() {
	local dev="$1"
	local postfix="$2"
	grep "$dev$postfix" $INTERRUPTS
}

set_affinity() {
	local dev="$1"
	local cpu_bound="$2"
	local queue_count="$3"
	local irq queue cpu cpus
	cpus="$(cpus_list_all | sed 's/.*/cpu&/')"
	# shellcheck disable=SC2046
	while read _ irq $cpus _ queue cpu _; do
		irq="${irq//:}"
		echo "  - $dev: irq $irq $queue -> $cpu"
		if [ "$DRYRUN" = 0 ]; then
			echo $cpu > /proc/irq/$irq/smp_affinity_list
		fi
	done
}

tune() {
	local dev="$1"
	local cpu_bound="$2"
	local postfix="$3"
	local queue_count
	echo "- Распределение прерываний $dev на сокете $cpu_bound"
	interrtupts "$dev" "$postfix" | cat -n > /tmp/interrupts.$dev
	queue_count="$(wc -l < /tmp/interrupts.$dev)"
	socket_cpu_list "$cpu_bound" "$queue_count" | cat -n > /tmp/cpu.$cpu_bound
	join /tmp/interrupts.$dev /tmp/cpu.$cpu_bound > /tmp/irq.$dev
	set_affinity "$dev" "$cpu_bound" "$queue_count" < /tmp/irq.$dev
}

igb() {
	tune "$1" "$2" "-TxRx-"
}

ixgbe() {
	tune "$1" "$2" "-TxRx-"
}

vmxnet3() {
	tune "$1" "$2" "-rxtx-"
}

# mlnx eth2-0 eth2-1 eth2-2 eth2-3

parse_params() {
	if [ "$1" == '--dry-run' ]; then
		DRYRUN=1
		shift
	fi
	if [ "$1" = '--help' ]; then
		usage
		exit 0
	fi
	dev="$1"
	cpu_bound="${2:-0}"
}

main() {
	local driver
	local postfix
	parse_params "$@"
	driver="$(ethtool -i "$dev" | grep driver | awk '{print $2}')"
	case $driver in
		igb | ixgbe | vmxnet3 )
			$driver $dev $cpu_bound
			;;
		* )
			echo "Driver $driver for dev $dev not supported"
			exit 1
			;;
	esac
}

main "$@"
