#!/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
	}
	shift
	DRYRUN=1
fi


usage() {
	echo "Usage: [PREFIX=-<rx|tx|RxTx|TxRX>-] $0 [--dry-run] [--test] <dev> [cpu socket]"
	echo
	echo "# 90% use-case:"
	echo "Example: $0 eth1"
	echo
	echo "# You may be like to have 'dry run' before really change affinities:"
	echo "Example: $0 --dry-run eth1"
	echo
	echo "# On multiprocessor machine it's better to explicitely define which socket will be used"
	echo "Example: $0 eth1 0"
	echo "Example: $0 eth2 1"
	echo
	echo "# Sometimes queue-prefix differs even for one driver. If so, you can redefine it:"
	echo "Example: PREFIX=-rx- $0 eth3"
	echo "Example: PREFIX=-tx- $0 eth3"
}

__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 ($postfix) на сокете $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
}

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}"
}

detect_postfixes() {
	grep -o "$1-.*" "$INTERRUPTS" | sed "s/$1//g" | tr -d '[0-9]' | sort -u
}

main() {
	local postfix
	parse_params "$@"
	for postfix in $(detect_postfixes $dev); do
		tune $dev $cpu_bound $postfix
	done
}

main "$@"
