#!/bin/sh
set -e

kill_child_jobs() {
    # From https://stackoverflow.com/a/23336595
    # Kills all child processes, not just jobs.
    pkill -P $$
}
cleanup() {
    kill_child_jobs
}

# From https://unix.stackexchange.com/a/240736
for sig in INT QUIT HUP TERM; do
  trap "
    cleanup
    trap - $sig EXIT
    kill -s $sig "'"$$"' "$sig"
done
trap cleanup EXIT


verbose=y
cpulimit_args="--lazy"
limit=""
pids=""
exes=""
paths=""
max_processes=100
max_depth=3
watch_interval=0
subprocess_watch_interval=0.5
while [ $# -gt 0 ]
do
    opt="$1"
    case "$opt" in
        *=*)
            arg="${1#*=}"
            opt="${1%%=*}"
            argshift=1
            ;;
        *)
            if [ "$#" -ge 2 ]
            then
                arg="$2"
            else
                arg=""
            fi
            argshift=2
            ;;
    esac

    case "$opt" in
        -h|--help)
            cat <<EOF
Usage: $0 [TARGET] [OPTIONS...] [-- PROGRAM]
   TARGET may be one or more of these (either TARGET or PROGRAM is required):
      -p, --pid=N        pid of a process
      -e, --exe=FILE     name of a executable program file
      -P, --path=PATH    absolute path name of a
                         executable program file
   OPTIONS for $0
          --max-depth=N  If 0, only target explicitly referenced processes.
                         Otherwise, target subprocesses up to N layers deep.
          --max-processes=N
                         Maximum number of processes to limit. After this
                         limit is reached, new processes will be limited
                         as old ones die.
          --watch-interval=INTERVAL
                         If 0 (default), targets will be selected at
                         setup. Otherwise, every INTERVAL (argument to
                         sleep(1)), search for more possible targets.
          --subprocess-watch-interval=INTERVAL
                         During setup, delay INTERVAL (argument to sleep(1))
                         between searches for more subprocesses to avoid
                         spending 100% CPU searching for targets.
          --             This is the final $0 option. All following
                         options are for another program we will launch.
      -h, --help         display this help and exit
   OPTIONS forwarded to CPUlimit
      -c, --cpu=N        override the detection of CPUs on the machine.
      -l, --limit=N      percentage of cpu allowed from 1 up.
                         Usually 1 - 800, but can be higher
                         on multi-core CPUs (mandatory)
      -q, --quiet        run in quiet mode (only print errors).
                         (Also suppresses messages from $0.)
      -k, --kill         kill processes going over their limit
                         instead of just throttling them.
      -r, --restore      Restore processes after they have
                         been killed. Works with the -k flag.
      -s, --signal=SIG   Send this signal to the watched process when cpulimit exits.
                         Signal should be specificed as a number or
                         SIGTERM, SIGCONT, SIGSTOP, etc. SIGCONT is the default.
      -v, --verbose      show control statistics
EOF
            exit 1
            ;;
        --max-depth)
            case $arg in
                ''|*[!0-9]*)
                    echo "Invalid max depth: $arg, must be a non-negative integer."
                    exit 5
                    ;;
                *)
                    max_depth=$arg
                    shift $argshift
                    ;;
            esac
            ;;
        --max-processes)
            case $arg in
                ''|*[!0-9]*)
                    echo "Invalid max processes: $arg, must be a positive integer."
                    exit 5
                    ;;
                *)
                    max_processes=$arg
                    shift $argshift
                    ;;
            esac
            ;;
        --watch-interval)
            watch_interval="$arg"
            shift $argshift
            ;;
        --subprocess-watch-interval)
            subprocess_watch_interval="$arg"
            shift $argshift
            ;;
        -p|--pid)
            if [ -z "$pids" ]
            then
                pids="$arg"
            else
                pids="$pids
$arg"
            fi
            shift $argshift
            ;;
        -e|--exe)
            if [ -z "$exes" ]
            then
                exes="$arg"
            else
                exes="$exes
$arg"
            fi
            shift $argshift
            ;;
        -P|--path)
            if [ -z "$paths" ]
            then
                paths="$arg"
            else
                paths="$paths
$arg"
            fi
            shift $argshift
            ;;
        -l|--limit)
            limit="$arg"
            cpulimit_args="$cpulimit_args --limit=$arg"
            shift $argshift
            ;;
        -c|--cpu)
            cpulimit_args="$cpulimit_args --cpu=$arg"
            shift $argshift
            ;;
        -v|--verbose)
            cpulimit_args="$cpulimit_args --verbose"
            shift
            ;;
        -q|--quiet)
            verbose=""
            cpulimit_args="$cpulimit_args --quiet"
            shift
            ;;
        -k|--kill)
            cpulimit_args="$cpulimit_args --kill"
            shift
            ;;
        -r|--restore)
            cpulimit_args="$cpulimit_args --restore"
            shift
            ;;
        -s|--signal)
            cpulimit_args="$cpulimit_args --signal=$arg"
            shift $argshift
            ;;
        --)
            shift
            break
            ;;
        *)
            echo "Unexpected argument: \"$1\""
            exit 3
            ;;
    esac
done

if [ -z "$limit" ]
then
    echo "Must specify a CPU percentage to limit to (-l/--limit)."
    exit 4
fi
if [ -z "$pids" ] && [ -z "$exes" ] && [ -z "$paths" ] && [ "$#" -eq 0 ]
then
    echo "Must specify at least one PID, executable file, path, or command line to limit."
    exit 5
fi

watched_pids=""
limit_pid() {
    if echo "$watched_pids" | grep --silent -F "$1"
    then
        return
    fi
    if [ "$(echo "$watched_pids" | wc -l)" -ge "$max_processes" ]
    then
        return
    fi

    if [ -n "$verbose" ]
    then
        if [ -n "$3" ]
        then
            echo "Limiting $3: "
        fi
        echo "cpulimit --pid=$1 $cpulimit_args"
    fi
    # shellcheck disable=SC2086  # $cpulimit_args really is the intended args.
    cpulimit --pid="$1" $cpulimit_args &
    cpulimit_pid=$!
    new_watched="$1:$2:$cpulimit_pid:$3"
    if [ -z "$watched_pids" ]
    then
        watched_pids="$new_watched"
    else
        watched_pids="$watched_pids
$new_watched"
    fi
}

limit_pids() {
    pids="$1"
    depth="$2"
    while read -r pid
    do
        # From https://stackoverflow.com/a/3951175
        case $pid in
            ''|*[!0-9]*)
                # PID is not a number
                ;;
            *)
                limit_pid "$pid" "$depth" "$3"
                ;;
        esac
    done <<EOF
$pids
EOF
}

limit_by_executable() {
    if [ -n "$exes" ]
    then
        while read -r exe
        do
            limit_pids "$(pgrep -x "$exe")" 0 "$exe"
        done <<EOF
$exes
EOF
    fi

    if [ -n "$paths" ]
    then
        while read -r path
        do
            limit_pids "$(pgrep -xf "$path")" 0 "$path"
        done <<EOF
$paths
EOF
    fi
}

limit_by_subprocess() {
    if [ -z "$watched_pids" ] || [ "$max_depth" -eq 0 ]
    then
        return
    fi

    while read -r watched
    do
        depth="$(echo "$watched" | cut -d: -f2)"
        if [ "$max_depth" -gt "$depth" ]
        then
            # Make sure the parent is still alive.
            if ps -p "$(echo "$watched" | cut -d: -f3)" >/dev/null
            then
                ppid="$(echo "$watched" | cut -d: -f1)"
                original="$(echo "$watched" | cut -d: -f4-)"
                limit_pids "$(pgrep -P "$ppid")" "$((depth + 1))" "child of $ppid ($original)"
            fi
        fi
    done <<EOF
$watched_pids
EOF
}

clean_dead_cpulimit() {
    if [ -z "$watched_pids" ]
    then
        return
    fi

    tmp="$(echo "$watched_pids" | while read -r watched
    do
        if ps -p "$(echo "$watched" | cut -d: -f3)" >/dev/null
        then
            echo "$watched"
        fi
    done)"
    watched_pids="$tmp"
}

if [ "$#" -gt 0 ]
then
    "$@" &
    limit_pid "$!" 0 "program run on command line: $*"
fi

limit_pids "$pids" 0
while true
do
    clean_dead_cpulimit
    if [ -z "$exes" ] && [ -z "$paths" ] && [ -z "$watched_pids" ]
    then
        # If there's nothing left to wait for, then exit.
        exit
    fi
    limit_by_executable
    if [ -z "$watched_pids" ]
    then
        num_watched_before=0
    else
        num_watched_before="$(echo "$watched_pids" | wc -l)"
    fi
    limit_by_subprocess
    if [ -z "$watched_pids" ]
    then
        num_watched_after=0
    else
        num_watched_after="$(echo "$watched_pids" | wc -l)"
    fi
    if [ "$num_watched_before" -eq "$num_watched_after" ]
    then
        if [ "$watch_interval" = "0" ]
        then
            if [ "$num_watched_after" -eq 0 ]
            then
                if [ -n "$verbose" ]
                then
                    echo "No processes found, exiting. Specify --watch-interval to continue scanning for processes."
                fi
                exit
            else
                if [ -n "$verbose" ]
                then
                    echo "Identified all processes to limit, waiting."
                fi
                wait
            fi
        else
            sleep "$watch_interval"
        fi
    else
        sleep "$subprocess_watch_interval"
    fi
done