#!/usr/bin/env bash
# memorycap -- Manage memory usage of processes
# Usage: memorycap [ARGS] COMMAND...
#
# Author: Dan Lawson <dan.lawson@bristol.ac.uk>
# Created: 6th Dec 2013
#
# Adapted from memusg by
# Author: Jaeho Shin <netj@sparcs.org>
function help {
    echo "Usage: memcap <options> <cmd>"
    echo "Cap the total memory used by a command"
    echo "Options:"
    echo "-m <val>      Kills the process if memory is detected to go above <val>"
    echo "-t <val>      Check memory every <val> seconds (default: 0.1)"
    echo "-h            Show this help"
    echo "-v            Verbose mode"
    echo "Set <val> to a negative number (default) to set no limit"
    echo "Prints memory usage in KB to stderr at the end of the run"
}
function checkfloat {
    val=`echo "$1" | grep -Eq '^[-+]?[0-9]+\.?[0-9]*$' && echo Match`
    if [ "$val" != "Match" ] ; then
	echo "Invalid float passed; received $1 but expected a floating point number"
	exit 1
    fi
}
function checkint {
    val=`echo "$1" | grep -Eq '^[-+]?[0-9]*$' && echo Match`
    if [ "$val" != "Match" ] ; then
	echo "Invalid int passed; received $1 but expected an integer"
	exit 1
    fi
}
function checkpos {
    if [ $1 -le 0 ]; then
	echo "Invalid number passed; received $1 but expected a positive number"
	exit 1
    fi
}

cap="-1"
verbose=""
interval=0.1

while getopts "a:m:t:hv" opt; do
  case $opt in
    m)
	  echo "-m"
      checkint $OPTARG
      cap=$OPTARG;;
    t) 
      checkfloat $OPTARG
      checkpos $OPTARG
      interval=$OPTARG;;
    h)
      help;exit 0;;
    v)
      verbose="T";;
    \?)
      exit 1
      ;;
    :)
      echo "Option -$OPTARG requires an argument." >&2
      exit 1
      ;;
  esac
done
shift $((OPTIND-1))

if [ $# -eq 0 ] ;then
    help
    exit 1
fi

##
set -um

# check input
[ $# -gt 0 ] || { sed -n '2,/^#$/ s/^# //p' <"$0"; exit 1; }
 
# TODO support more options: peak, footprint, sampling rate, etc.
 
pgid=`ps -o pgid= $$`
# make sure we're in a separate process group
if [ $pgid = $(ps -o pgid= $(ps -o ppid= $$)) ]; then
cmd=
set -- "$0" "$@"
for a; do cmd+="'${a//"'"/"'\\''"}' "; done
exec bash -i -c "$cmd"
fi
 
# detect operating system and prepare measurement
case `uname` in
Darwin|*BSD) sizes() { /bin/ps -o rss= -g $1; } ;;
Linux) sizes() { /bin/ps -o rss= -$1; } ;;
*) echo "`uname`: unsupported operating system" >&2; exit 2 ;;
esac

date1=$(date +"%s")
echo date 1>&2
 
# monitor the memory usage in the background.
(
peak=0
while sizes=`sizes $pgid`
do
    set -- $sizes
    sample=$((${@/#/+}))
## Print anything 
    if [ "$verbose" == "T" ] ; then
	if [ $sample -gt $peak ] ;then
	    echo "Updating peak to $peak"
	fi
    fi
## Update the peak memory usage
    let peak="sample > peak ? sample : peak"
## Check whether to kill the program
    if [ $peak -gt $cap ] ;then
	if [ $cap -gt 0 ] ;then
	    echo "KILLING PROCESS $pgid DUE TO MEMORY RESTRICTIONS: cap $cap < usage $peak"
	    kill $pgid
#	    sleep 10
#	    kill -9 $pgid
	fi
    fi
    sleep $interval
done
date2=$(date +"%s")
date 1>&2
diff=$(($date2-$date1))
echo "Duration in Seconds: $diff" 1>&2
echo "memusg: peak = $peak" >&2
) &
monpid=$!
 
 
# run the given command
if [ "$verbose" == "T" ] ; then
    echo "Running command $@"
    echo "with cap $cap"
fi
exec "$@"
