#!/bin/bash
VERBOSE=y

command=firefox
memory_ratio=0.5
control_group=$command

as_root=sudo

total_memory() { free -b | (read line; read Mem total _; echo $total); }
out() { $as_root tee "$@" >/dev/null; }

die() { printf 'Error: %s\n' "$*" >&2; exit 1; }
verbose() { [ ! "$VERBOSE" ] || printf '%s\n' "$*" >&2; }

total_memory=$(total_memory) || die "could not ascertain total system memory; probable bug"
memory_limit_in_bytes=$(bc -q <<< "$total_memory * $memory_ratio / 1" ) || die "process 'bc' failed.  Is it installed?"

init_control_group()
{
    control_group_dir=/sys/fs/cgroup/"$1"/"$control_group"
    [ -d "$control_group_dir" ] && return
    $as_root mkdir "$control_group_dir" || die "mkdir failed"
    set_memory_limit $memory_limit_in_bytes || die "set_memory_limit failed"
}

get_cpu_limit()
{
    false
}

get_memory_limit()
{
    cat "$control_group_dir"/memory.limit_in_bytes
}

set_memory_limit()
{
    prev=$(get_memory_limit)
    if [ "$1" -ne "$prev" ]; then
        out "$control_group_dir"/memory.limit_in_bytes <<< $1
        verbose "changed memory limit ($prev -> $1)"
    else
        verbose "memory limit unchanged ($prev)"
    fi
}

join_control_group()
{
    out "$control_group_dir"/cgroup.procs <<< $$
}

get_cgroup()
{
    local key="$1" pid="${2:-$$}" line val
    while read line; do
        line=${line#*:}
        val=${line#$key:}
        [ "$val" != "$line" ] || continue
        echo "$val"
        return
    done < /proc/"$pid"/cgroup
    false
}

launch_command()
{
    init_control_group memory               || die "init_control_group failed"
    cg=$(get_cgroup memory)                 || die "failed to ascertain current control group"
#   [ "$cg" = / ]                           || die "current control group should be '/', but is instead '$cg'"
    join_control_group                      || die "join_control_group failed"
    exec $command "$@"
}

pretty_show_memory_limit()
{
    local memory_limit="$(get_memory_limit)"
    printf "current memory limit: %d bytes (%d%%)\n" \
           "$memory_limit" \
           "$((memory_limit * 100 / total_memory))"
}

recpu_control_group()
{
    init_control_group cpu
    case "$1" in
        [0-9]*%) set_cpu_limit $((total_cpu * "${1%\%}" / 100)) ;;
        [0-9]*) set_cpu_limit "$1" ;;
        [-+*/][0-9]*) set_cpu_limit $(( $(get_cpu_limit) "$1" )) ;;
    esac
}

resize_control_group()
{
    init_control_group memory
    case "$1" in
        '') pretty_show_memory_limit >&2 ;;
        [0-9]*%) set_memory_limit $((total_memory * "${1%\%}" / 100)) ;;
        [0-9]*) set_memory_limit "$1" ;;
        [-+*/][0-9]*) set_memory_limit $(( $(get_memory_limit) "$1" )) ;;
    esac
}

usage()
{
    cat >&2 <<EOF
usage: ${0##*/} [--resize|--resize <bytes>|--resize=<bytes>] [--] [pass-through-options] [...]
       ${0##*/} --help
EOF
}

case "$1" in
    --) shift; launch_command "$@" ;;
    --resize) resize_control_group "$2" ;;
    --resize=*) resize_control_group "${1#--resize=}" ;;
    --cpu=*) recpu_control_group "${1#--cpu=}" ;;
    --*) usage; [ "$1" = '--help' -o "$1" = '--usage' ] ;;
    *) launch_command "$@" ;;
esac