#!/bin/bash

WEB_CONTENT_OOM_ADJ=500
FIREFIX_REPEAT_INTERVAL=5

ploop()
{
    local p
    for p in /proc/[0-9]*
    do
        "$@"
    done
}

match_parent()
{
    # This isn't a file to which we will write (we write to the child's
    # oom_score_adj), but this ignores processes we don't own.
    2>/dev/null [ -w "$p"/oom_score_adj ] || return

    local stat
    2>/dev/null read stat < "$p"/stat || return
    case "$stat" in
        *") "?" $1 "*) echo "${p##*/}" ;;
    esac
}

match_comm()
{
    local comm
    2>/dev/null read comm < "$p"/comm || return
    [ "$comm" = "$1" ] && echo "${p##*/}"
}

firefix()
{
    parent=$1
    read parent_oom_score < /proc/$parent/oom_score
    read parent_oom_score_adj < /proc/$parent/oom_score_adj

    for child in $(ploop match_parent $parent)
    do
        2>/dev/null read comm < /proc/$child/comm || continue
        [ "$comm" = 'Web Content' ] || continue

        read oom_score < /proc/$child/oom_score
        read oom_score_adj < /proc/$child/oom_score_adj

        want_adj=$((parent_oom_score_adj + ${WEB_CONTENT_OOM_ADJ:-500}))
        if [ "$want_adj" -gt "$oom_score_adj" ]
        then
            printf 'Setting oom_score_adj for pid=%d to %d (from %d)\n' $child $want_adj $oom_score_adj >&2
            printf '%d\n' "$want_adj" > /proc/$child/oom_score_adj
        fi
    done
}

firefix_all()
{
    for parent in $(ploop match_comm firefox-bin)
    do
        firefix $parent &
    done
    wait
}

firefix_all_forever()
{
    while true
    do
        firefix_all
        sleep ${FIREFIX_REPEAT_INTERVAL:-5}
    done
}

unit_file()
{
    cat <<EOF
[Unit]
Description=$1
[Service]
ExecStart=$2
[Install]
WantedBy=default.target
EOF
}

install_self()
{
    unit_file_name=$1
    unit_executable=$2
    unit_args=$3
    unit_desc=$4

    [ -e "$unit_executable" ] || return

    if [ "$(id -u)" = 0 ]
    then
        service_dir=/etc/systemd/system
        systemctl=systemctl
        instdir=/usr/local/bin
    else
        service_dir=$HOME/.config/systemd/user
        systemctl='systemctl --user'
        instdir=$HOME/.local/bin
    fi
    unit_executable_installed=$instdir/${unit_executable##*/}

    [ "$unit_executable_installed" -ef "$unit_executable" ] ||
        install -D -t "$instdir" "$unit_executable" || return

    [ -d "$service_dir" ] || mkdir -p "$service_dir" || return

    unit_file=${service_dir}/${unit_file_name}.service
    unit_file "$unit_desc" "$unit_executable_installed $unit_args" > "$unit_file"

    $systemctl daemon-reload
    $systemctl enable "$unit_file_name"
    $systemctl restart "$unit_file_name"
    $systemctl status "$unit_file_name"
}

usage()
{
    echo "Usage: $0 <install|once|forever>" >&2
}

enable -f /usr/lib/bash/sleep sleep 2>/dev/null || true

case "$*" in
    forever) firefix_all_forever ;;
    install) install_self firefixer "$0" forever 'Firefixer - adjust firefox OOM scores';;
    once) firefix_all ;;
    -h|--help) usage; exit ;;
    *) usage; exit 1 ;;
esac