#!/bin/bash
# ==============================================================
# CPU limit daemon - set PID's max. percentage CPU consumptions
# ==============================================================

# Variables
CPU_LIMIT=32

LOAD_AVERAGE_LIMIT=2 # Уровень LA после которого демон считает что система перегружена (установите 0, если не знаете что это)

DAEMON_INTERVAL=2   # Daemon check interval in seconds

INTERVALS_BEFORE_THAW=8 # Количество интервалов без System overload или выставления лимитов, пропускаемых от снятия всех лимитов

BLACK_PROCESSES_LIST="vivaldi|opera|megasync"   # Черный список приложений, которые будут ограничиваться даже если система не нагружена
WHITE_PROCESSES_LIST=""   # Белый список приложений, которые не буду ограничиваться даже если система перегружена

################################################################################
if [[ $LOAD_AVERAGE_LIMIT -eq 0 ]] ; then
    LOAD_AVERAGE_LIMIT=`grep -c '^processor' /proc/cpuinfo`
fi

if [[ $WHITE_PIDS_COMMAND != '' ]] ; then
    WHITE_PIDS_COMMAND="LC_ALL=C top -b -n1 -c | gawk 'NR>6' | grep -E -v '$WHITE_PROCESSES_LIST' | gawk '\$9>CPU_LIMIT {print \$1}' CPU_LIMIT=$CPU_LIMIT";
else
    WHITE_PIDS_COMMAND="LC_ALL=C top -b -n1 -c | gawk 'NR>6' | gawk '\$9>CPU_LIMIT {print \$1}' CPU_LIMIT=$CPU_LIMIT";
fi
BLACK_PIDS_COMMAND="LC_ALL=C top -b -n1 -c | grep -E '$BLACK_PROCESSES_LIST' | gawk '\$9>CPU_LIMIT {print \$1}' CPU_LIMIT=$CPU_LIMIT"

STATE_COMMAND="cat /proc/loadavg | gawk '\$1>LOAD_AVERAGE_LIMIT {print overload}' LOAD_AVERAGE_LIMIT=$LOAD_AVERAGE_LIMIT overload='overload'" ;

BEFORE_THAW=-1;
# Search and limit violating PIDs
################################################################################
while sleep $DAEMON_INTERVAL
do
echo '==============================================================================='
    LOAD_AVERAGE_STATE=$(eval "$STATE_COMMAND");
    if [[ $LOAD_AVERAGE_STATE == 'overload' ]] ; then
        NEW_PIDS=$(eval "$WHITE_PIDS_COMMAND");
        BEFORE_THAW=$INTERVALS_BEFORE_THAW; # перегрузка - полностью откладываем оттепель
        echo "System overload!";
    else
        if [[ $BEFORE_THAW -eq 0 ]] ; then
            killall cpulimit;
            BEFORE_THAW=-1;
            echo 'Thaw';
            sleep $DAEMON_INTERVAL;
        else
            BEFORE_THAW=$(($BEFORE_THAW - 1));
            echo "Intervals before thaw: $BEFORE_THAW"
        fi
        NEW_PIDS=$(eval "$BLACK_PIDS_COMMAND");
    fi

    if [[ $NEW_PIDS != '' ]] ; then # есть процессы нарушители
        NEW_PIDS=$(echo $NEW_PIDS)
#echo ''
#echo 'New pids:'
#LC_ALL=C top -b -n1 -c -p ${NEW_PIDS//[[:space:]]/" -p "} | gawk 'NR>6'

        LIMITED_PIDS=$(ps -eo args | gawk '($1=="cpulimit") {print $3}' CPU_LIMIT=$CPU_LIMIT) # PID'ы в отношении которых уже действую ограничения
        PIDS4LIMIT=$(comm -23 <(echo "$NEW_PIDS" | sort -u) <(echo "$LIMITED_PIDS" | sort -u) | grep -v '^$')

        if [[ $PIDS4LIMIT != '' ]] ; then
            BEFORE_THAW=$(($INTERVALS_BEFORE_THAW - 1)); # на один инервал меньше
            for i in $PIDS4LIMIT # перебираем процессы которые следуюет ограничить и выставляем для них предварительное ограничение
            do
                cpulimit -p $i -l $CPU_LIMIT -z & # Limit new violating processes
            done
        fi

    fi

LIMITED_PIDS=$(ps -eo args | gawk '($1=="cpulimit") {print $3}' CPU_LIMIT=$CPU_LIMIT) # PID'ы в отношении которых уже действую ограничения
if [[ $LIMITED_PIDS != '' ]] ; then
    echo ''
    echo 'Current limited pids:'
    LC_ALL=C top -b -n1 -c -p ${LIMITED_PIDS//[[:space:]]/" -p "} | gawk 'NR>6'
fi
done