#!/bin/bash
set -eu
echo "$0 $@ [$$] START" >&2
if [ "${1:-}" = "--help" ]; then
	echo 'Info: Ограничение cpu affinity по количеству процессов и cpu usage'
	echo 'Usage:'
	echo 'Example:'
	exit 0
fi

. /opt/fox_node/__node_config

declare NODE_TYPE RUN MAX_PROCESS_COUNT NODE_CPUS MAX_CPU_USAGE_BY_CTID MAX_NODE_CPU_USAGE
declare MAX_NODE_LA NODE_CPU_MASK_DEFAULT CPU_MASK_TO_CHANGE NEED_CHANGE_AFFINITY
declare ALL_CHANGED_PIDS ALARM_VM_NAMES ALARM_TS
declare -A VM_PIDS
declare -A VM_PIDS_TO_CLEAR_AFFINITY
declare -A VM_PIDS_TO_SET_AFFINITY
declare -A VM_CPU_MASK
declare -A VM_NAMES

# Для корректного разделителя точка при ручном запуске
export LANG="en_US.UTF-8"

# Пока тестируем без реального ограничения процессов, не будет выполняться taskset
RUN="echo"
MAX_PROCESS_COUNT=100
NODE_CPUS="$(grep -c "processor" "/proc/cpuinfo")"
# Т.к. cpu usage проверяем через top, то все значения доводим до целого числа *10
MAX_CPU_USAGE_BY_CTID=900
MAX_NODE_CPU_USAGE=900
MAX_NODE_LA=18
NODE_CPU_MASK_DEFAULT="0-$((NODE_CPUS-1))"
CPU_MASK_TO_CHANGE="$((NODE_CPUS-1))"
NEED_CHANGE_AFFINITY=FALSE
ALL_CHANGED_PIDS="/var/lib/process_cpu_affinity_balance/changed_pids"
ALARM_VM_NAMES=""
ALARM_TS="/var/lib/process_cpu_affinity_balance/alarm_ts"

_get_int_val() {
	set -e
	local val="$1" ret
	ret="$(awk "BEGIN {print ($val)}")"
	echo "${ret%.*}"
	return 0
}

_get_node_cpu_usage() {
	set -e
	local top_count node_cpu_usage_total node_cpu_usage_list cpu_usage_current
	top_count=3
	node_cpu_usage_total=0
	# Считаем общий cpu usage ноды
	# %Cpu(s): 26.2 us.  2.1 sy.  0.0 ni. 71.7 id.  0.0 wa.  0.0 hi.  0.0 si.  0.0 st
	# или
	# %Cpu(s):100.0 us.  0.0 sy.  0.0 ni. 0.0 id.  0.0 wa.  0.0 hi.  0.0 si.  0.0 st
	node_cpu_usage_list=$(top -bn $top_count -d 1 | grep -oE "Cpu\(s\):.* us")
	node_cpu_usage_list=$(echo "$node_cpu_usage_list" | cut -d':' -f2 | cut -d'u' -f1)
	for cpu_usage_current in $node_cpu_usage_list; do
		node_cpu_usage_total=$(_get_int_val "$cpu_usage_current * 10 + $node_cpu_usage_total")
	done
	_get_int_val "$node_cpu_usage_total / $top_count"
	return 0
}

_get_node_loadaverage() {
	set -e
	local la_1_min la_5_min la_15_min t
	read -r la_1_min la_5_min la_15_min t < /proc/loadavg
	case "$1" in
	1)
		_get_int_val "$la_1_min"
		;;
	5)
		_get_int_val "$la_5_min"
		;;
	15)
		_get_int_val "$la_15_min"
		;;
	*)
		echo "Не поддерживается параметр load average $1, только 1,5,15"
		exit 255
		;;
	esac
	return 0
}

_set_vm_pids_to_clear_affinity() {
	local ctid="$1" pid="$2" t cpus_allowed_list
	read -r t cpus_allowed_list <<< "$(grep -s "Cpus_allowed_list:" "/proc/${pid}/status")"
	if [ -n "$cpus_allowed_list" ] && [ "$cpus_allowed_list" != "${VM_CPU_MASK["$ctid"]}" ]; then
		# Ограничили процесс именно мы
		if grep -q "$pid" "$ALL_CHANGED_PIDS"; then
			# Все процессы с изменным cpu affinity для вм
			VM_PIDS_TO_CLEAR_AFFINITY["$ctid"]="${VM_PIDS_TO_CLEAR_AFFINITY["$ctid"]:-} $pid"
		fi
	fi
	return 0
}

check_node_usage() {
	local node_cpu_usage node_la_5_min
	node_cpu_usage="$(_get_node_cpu_usage)"
	node_la_5_min="$(_get_node_loadaverage 5)"
	NEED_CHANGE_AFFINITY=FALSE
	# ТODO, проверка на disk i/o
	if [ "$node_cpu_usage" -ge "$MAX_NODE_CPU_USAGE" ]; then
		echo "node_cpu_usage = $node_cpu_usage > $MAX_NODE_CPU_USAGE"
		echo "top -bn 1 | head -n 20"
		top -bn 1 | head -n 20
		NEED_CHANGE_AFFINITY=TRUE
	fi
	if [ "$node_la_5_min" -ge "$MAX_NODE_LA" ]; then
		echo "node_la_5_min = $node_la_5_min > $MAX_NODE_LA"
		echo "node_cpus_count = $NODE_CPUS"
		echo "cat /proc/loadavg"
		cat /proc/loadavg
		echo "top -bn 1 | head -n 20"
		top -bn 1 | head -n 20
		NEED_CHANGE_AFFINITY=TRUE
	fi
	return 0
}

prepare_params() {
	local ctid vm_name cpumask
	if [ ! -f "$ALL_CHANGED_PIDS" ]; then
		mkdir -p "${ALL_CHANGED_PIDS%/*}"
		touch "$ALL_CHANGED_PIDS"
	fi
	if [ ! -f "$ALARM_TS" ]; then
		mkdir -p "${ALARM_TS%/*}"
		touch "$ALARM_TS"
	fi

	while read -r ctid vm_name cpumask; do
		VM_NAMES["$ctid"]="$vm_name"
		# Маска по-умолчанию для вм, если у кого-то ограничили вм через
		# vzctl set vps123456 --set--cpumask=1,2, то берем её маску.
		# Эта маска будет использоваться для снятия ограничений с процессов.
		if [ "$cpumask" = "-" ]; then
			VM_CPU_MASK["$ctid"]="$NODE_CPU_MASK_DEFAULT"
		else
			VM_CPU_MASK["$ctid"]="$cpumask"
		fi
		while read -r pid; do
			# Все процессы вм
			VM_PIDS["$ctid"]="${VM_PIDS["$ctid"]:-} $pid"
			_set_vm_pids_to_clear_affinity "$ctid" "$pid"
		done < "/sys/fs/cgroup/pids/machine.slice/${ctid}/tasks"
	done <<< "$(vzlist -H -o ctid,name,cpumask)"
	# Ядра на которые будем ставить cpu affinity в зависимости от ноды
	if [ "$NODE_CPUS" -ge 12 ]; then
		CPU_MASK_TO_CHANGE="$((NODE_CPUS-3))-$((NODE_CPUS-1))"
	elif [ "$NODE_CPUS" -ge 8 ]; then
		CPU_MASK_TO_CHANGE="$((NODE_CPUS-2))-$((NODE_CPUS-1))"
	fi
	return 0
}

check_ps() {
	local ctid vzps_list process_count_list count cmd n pid
	for ctid in ${!VM_PIDS[*]}; do
		vzps_list="$(vzps -E "$ctid")"
		# Посчитаем сколько одноименных процессов запущено у вм
		process_count_list="$(echo "$vzps_list" | awk '{print $5}' | sort | uniq -c)"
		while read -r count cmd; do
			if [ "$count" -ge "$MAX_PROCESS_COUNT" ]; then
				echo "У ${VM_NAMES["$ctid"]} количество процессов $cmd $count > $MAX_PROCESS_COUNT"
				while read -r n pid n; do
					# Помечаем вмку и конкретные pid для ограничения
					VM_PIDS_TO_SET_AFFINITY["$ctid"]="${VM_PIDS_TO_SET_AFFINITY["$ctid"]:-} $pid"
				done <<< "$(grep "$cmd" <<< "$vzps_list")"
			fi
		done <<< "$process_count_list"
	done
	return 0
}

check_top() {
	local top_count ctid total_cpu_usage pid n cpu cmd
	declare -A total_cpu_usage_by_cmd pids_by_cmd
	top_count=3
	# Сделаем несколько (top_count) замеров top раз в 1 сек
	top -bn "$top_count" -d 1 > "/tmp/process_cpu_affinity_balance.$$"
	for ctid in ${!VM_PIDS[*]}; do
		total_cpu_usage=0
		total_cpu_usage_by_cmd=()
		pids_by_cmd=()
		# top - 10:18:09 up 125 days, 20:28,  3 users,  load average: 0.12, 0.07, 0.06
		# Tasks: 597 total,   1 running, 596 sleeping,   0 stopped,   0 zombie
		# %Cpu(s):  0.1 us,  0.0 sy,  0.0 ni, 99.9 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
		# KiB Mem : 65705752 total, 43218172 free, 11116420 used, 11371160 buff/cache
		# KiB Swap: 15769376 total, 15088172 free,   681204 used. 49162244 avail Mem
		#
		#    PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND
		while read -r pid n n n n n n n cpu n n cmd n; do
			[ "$pid" != "${pid//[!0-9]/}" ] && continue  # только цифры
			[ -z "$pid" ] && continue
			if [[ " ${VM_PIDS["$ctid"]} " == *" $pid "* ]]; then
				cpu_usage=$(_get_int_val "$cpu * 10")  # для удобного суммирования, т.к. будет int
				if [ -z "${total_cpu_usage_by_cmd["$cmd"]:-}" ]; then
					total_cpu_usage_by_cmd["$cmd"]=0
				fi
				# Будем считать cpu usage для одноименных процессов и общую для вм
				total_cpu_usage_by_cmd["$cmd"]=$((total_cpu_usage_by_cmd["$cmd"] + cpu_usage))
				total_cpu_usage=$((total_cpu_usage + cpu_usage))
				# Т.к. несколько повторений top, то отфильтруем лишние pid
				# и не берем те, которые уже ограничены
				if [[ " ${pids_by_cmd["$cmd"]:-} " != *" $pid "* ]] \
					&& [[ " ${VM_PIDS_TO_CLEAR_AFFINITY["$ctid"]:-} " != *" $pid "* ]]; then
					pids_by_cmd["$cmd"]="${pids_by_cmd["$cmd"]:-} $pid"
				fi
			fi
		done < "/tmp/process_cpu_affinity_balance.$$"
		if [ "$total_cpu_usage" -ge "$((MAX_CPU_USAGE_BY_CTID * top_count))" ]; then
			echo "У ${VM_NAMES["$ctid"]} cpu_usage > $MAX_CPU_USAGE_BY_CTID"
			# Отсортируем процессы по cpu usage
			set -o pipefail
			for cmd in "${!total_cpu_usage_by_cmd[@]}"; do
				echo "$cmd" "${total_cpu_usage_by_cmd["$cmd"]}"
			done | sort -rn -k2 > "/tmp/process_cpu_affinity_balance_$ctid.$$"
			set +o pipefail
			echo "top 10 процессов ${VM_NAMES["$ctid"]} by cpu_usage"
			head -n 10 "/tmp/process_cpu_affinity_balance_$ctid.$$"
			# Для ограничения берем процесс с самым большим cpu usage
			# если этот процесс уже ограничен (нет pid), то берем следующий
			while read -r cmd n; do
				if [ -n "${pids_by_cmd["$cmd"]:-}" ]; then
					# Здесь возможно дублирование pid после check_ps, но не страшно
					VM_PIDS_TO_SET_AFFINITY["$ctid"]="${VM_PIDS_TO_SET_AFFINITY["$ctid"]:-} ${pids_by_cmd["$cmd"]}"
					echo "Для ограничения выбран процесс $cmd"
					break
				fi
			done < "/tmp/process_cpu_affinity_balance_$ctid.$$"
			rm -f "/tmp/process_cpu_affinity_balance_$ctid.$$"
		fi
	done
	rm -f "/tmp/process_cpu_affinity_balance.$$"
	return 0
}

change_cpu_affinity() {
	local ctid pid
	# Будем ограничивать вм пока не упадет нагрузка для ноды ниже определенной
	# или не закончаться вмки
	for ctid in "${!VM_PIDS_TO_SET_AFFINITY[@]}"; do
		if [ "$NEED_CHANGE_AFFINITY" = TRUE ]; then
			echo "Ограничена вм ${VM_NAMES["$ctid"]}"
			for pid in ${VM_PIDS_TO_SET_AFFINITY[$ctid]}; do
				$RUN taskset -c -pa "$CPU_MASK_TO_CHANGE" "$pid"
				echo "$pid" >> "$ALL_CHANGED_PIDS"
			done
			ALARM_VM_NAMES="${ALARM_VM_NAMES} ${VM_NAMES["$ctid"]}"
		else
			break
		fi
		check_node_usage
	done
	return 0
}

clear_cpu_affinity() {
	local ctid pid
	# Будем снимать ограничения пока не повысится нагрузка для ноды выше определенной
	# или не закончаться вмки
	for ctid in "${!VM_PIDS_TO_CLEAR_AFFINITY[@]}"; do
		if [ "$NEED_CHANGE_AFFINITY" = FALSE ]; then
			if [[ " ${!VM_PIDS_TO_SET_AFFINITY[*]:-} " != *" $ctid "* ]]; then
				echo "Снятие ограничений с вм ${VM_NAMES["$ctid"]}"
				# Возможно стоит снимать ограничения не со всех процессов, а пачками xx штук
				for pid in ${VM_PIDS_TO_CLEAR_AFFINITY[$ctid]}; do
					$RUN taskset -c -pa "${VM_CPU_MASK["$ctid"]}" "$pid"
					sed -i "/$pid/d" "$ALL_CHANGED_PIDS"
				done
			fi
		else
			break
		fi
		check_node_usage
	done
	return 0
}

check_dead_pids() {
	sort -u "$ALL_CHANGED_PIDS" > "$ALL_CHANGED_PIDS.tmp"
	while read -r pid; do
		if [ ! -d "/proc/$pid" ]; then
			sed -i "/$pid/d" "$ALL_CHANGED_PIDS"
		fi
	done < "$ALL_CHANGED_PIDS.tmp"
	rm -f "$ALL_CHANGED_PIDS.tmp"
	return 0
}

alarm() {
	local cur_ts last_ts=0 diff_ts vm_name
	cur_ts="$(date +%s)"
	if [ "$ALARM_VM_NAMES" != "" ]; then
		for vm_name in $ALARM_VM_NAMES; do
			last_ts="$(/opt/fox_utils/crab_conf get "$vm_name" "$ALARM_TS")"
			diff_ts="$((cur_ts-last_ts))"
			if [ "$diff_ts" -ge "$((60*60*24*7))" ]; then
				/opt/fox_utils/fox_alarm "Ограничена вм $vm_name"
				/opt/fox_utils/crab_conf set "$vm_name" "$cur_ts" "$ALARM_TS"
				echo "Send alarm /opt/fox_utils/fox_alarm 'Ограничена вм $vm_name'"
			else
				echo "Skip alarm, diff_ts = $diff_ts < $((60*60*24*7))"
			fi
		done
	fi
	return 0
}

main() {
	echo "Started $(date +"%Y-%m-%d %H:%M:%S")"
	if [[ " ${NODE_TYPE:-} " != *" vz7 "* ]]; then
		echo "Поддерживаются только ноды vz7"
		exit 0
	fi

	check_node_usage
	prepare_params
	if [ "$NEED_CHANGE_AFFINITY" = TRUE ]; then
		check_ps
		check_top
		change_cpu_affinity
	else
		# В будущем нужно защищаться от бесконечного вкл/выкл т.к.taskset повлияет на нагрузку ноды.
		# Мб вести в будущем "статистику банов" в разрезе виртуалок:
		# - первый бан на 1 минуту
		# - если после первого разбана вернулась нагрузка - баним на 5 мин
		# - итд
		clear_cpu_affinity
	fi
	check_dead_pids
	alarm
	return 0
}

main "$@"

echo "$0 $@ [$$] SUCCESS" >&2
exit 0