#!/usr/bin/bash
# Offline memory until the specified limit is reached.

FREE_FIELD=2
while [ $# -gt 0 ]; do
	case $1 in
	--limit)
		LIMIT=$2
		shift 2
		;;
	--limit-free)
		# By default, free until available memory matches the limit
		# --limit-free means limiting until free memory hits the limit
		# after dropping cache
		echo 3 > /proc/sys/vm/drop_caches
		FREE_FIELD=4
		shift
		;;
	*)
		echo WARNING: Unrecognised option $1
		shift
		;;
	esac
done

if ! [[ $LIMIT =~ ^[0-9]+$ ]] ; then
	echo Memory limit must be specified in bytes $LIMIT
	echo Usage: limit-memory --limit [bytes]
	exit -1
fi

if [ "`whoami`" != "root" ]; then
	echo limit-memory requires root
	exit -1
fi

SCRIPT_ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
MEMORY_BANKS=`$SCRIPT_ROOT/enum-memory.pl`
if [ "$MEMORY_BANKS" = "" ]; then
	echo Unable to enumerate available memory
	exit -1
fi

check_too_limited() {
	if [ $MEMTOTAL_BYTES -le $((LIMIT*4/5)) ]; then
		echo WARNING: Available memory is below 80% of the requested limit
		echo WARNING: Offlining was too aggressive and it may be required
		echo WARNING: to boot with mem=$((MEMTOTAL_BYTES/1048576))M instead.
	fi
}

MEMTOTAL_BYTES=`free -b | grep Mem: | awk "{print \\$\$FREE_FIELD}"`
if [ $MEMTOTAL_BYTES -le $LIMIT ]; then
	echo Memory limit already reached
	check_too_limited
	exit 0
fi

for BANK in $MEMORY_BANKS; do
	if [ ! -e /sys/devices/system/memory/memory$BANK/online ]; then
		echo Unable to detect state of memory bank $BANK
		continue
	fi

	if [ ! -e /sys/devices/system/memory/memory$BANK/state ]; then
		echo Unable to hotplug memory bank $BANK
		continue
	fi

	if [ ! -e /sys/devices/system/memory/memory$BANK/removable ]; then
		echo Bank $BANK is not removable
		continue
	fi
	REMOVABLE=`cat /sys/devices/system/memory/memory$BANK/removable`
	if [ $REMOVABLE -ne 1 ]; then
		echo Bank $BANK is not removable
		continue
	fi

	echo -n offline > /sys/devices/system/memory/memory$BANK/state 2> /dev/null &
	OFFLINE_PID=$!
	ATTEMPT=0
	STATE=`cat /sys/devices/system/memory/memory$BANK/online`
	while [ $STATE -ne 0 ]; do
		if [ $ATTEMPT -eq 0 ]; then
			echo -n Waiting on bank $BANK for 3 seconds
		fi
		ATTEMPT=$((ATTEMPT+1))
		if [ $ATTEMPT -ge 3 ]; then
			echo Bank $BANK has not reached offline state
			ps -p $OFFLINE_PID &>/dev/null
			if [ $? -eq 0 ]; then
				kill -9 $OFFLINE_PID
			fi
			sleep 1
			echo -n online > /sys/devices/system/memory/memory$BANK/state &>/dev/null
			break
		fi
		sleep 1
		echo -n .
		STATE=`cat /sys/devices/system/memory/memory$BANK/online`
	done
	if [ $STATE -eq 0 ]; then
		echo Offlined bank $BANK
		MEMTOTAL_BYTES=`free -b | grep Mem: | awk "{print \\$\$FREE_FIELD}"`
		if [ $MEMTOTAL_BYTES -le $LIMIT ]; then
			break
		fi
	fi
done

MEMTOTAL_BYTES=`free -b | grep Mem: | awk "{print \\$\$FREE_FIELD}"`
if [ $MEMTOTAL_BYTES -gt $LIMIT ]; then
	free -m
	echo FAILED: Unable to limit memory, suggest booting with mem=$((LIMIT/1048576))M
	exit -1
fi

check_too_limited
echo Memory limit reached
exit 0
