#!/usr/bin/env bash
# http://brainsuckerna.blogspot.com/2016/05/limiting-mongodb-memory-usage-with.html
# (c) Miadzvedz Mikalai, 2016. brainsucker.na at gmail.com
# init scripts (c) frank2, see http://frank2.net/cgroups-ubuntu-14-04/ for description and details
# more info (no init scripts): using cgroups to limit DB memory https://www.percona.com/blog/2015/07/01/using-cgroups-to-limit-mysql-and-mongodb-memory-usage/
#
# execute from shell:
# sudo bash -c 'curl -o- http://brains.by/misc/mongodb_memory_limit_ubuntu1404.sh | bash'

{ # this ensures the entire script is downloaded #

echo "------------------------------------------------------"

if (( $EUID != 0 )); then
    echo "Please run as root"
    exit
fi

echo "Limit MongoDB Memory for Ubuntu 14.04 (with cgroups)"
echo "------------------------------------------------------"
echo "This script will OVERWRITE your cgroups configuration and init files."

read -p "Continue (y/n)?" -n 1 -r choice < /dev/tty
case "$choice" in 
  y|Y ) echo "yes";;
  n|N ) echo "no"; exit;;
  * ) echo "invalid"; exit;;
esac

echo "------------------------------------------------------"
echo "Please specify memory limit (in bytes) to use, like 5G, 512M or 100K."
read -p "Memory limit (default 256M)?"  -r LIMIT < /dev/tty

if [ -z "$LIMIT" ] 
then
	LIMIT="256M"
else
	if [[ ! "$LIMIT" =~ ^([0-9]+[GMKgmk]?)$ ]]  
	then
		echo "Couldn't parse Memory limit value (must be in bytes, or like 5G, 5M, 5K): $LIMIT"	
		exit
	fi
fi

echo "Using memory limit '$LIMIT'"
echo "------------------------------------------------------"

# install cgroups
#apt-get update
apt-get install cgroup-bin


echo "------------------------------------------------------"

echo "Creating /etc/cgconfig.conf"
cat > /etc/cgconfig.conf <<EOF
group DBLimitedGroup {
  memory {
    memory.limit_in_bytes = "$LIMIT";
  }
}
EOF

echo "Creating /etc/cgrules.conf"
cat >> /etc/cgrules.conf <<EOF
*:mongod	memory		DBLimitedGroup
EOF

echo "Creating /etc/init/cgroup-lite.conf"
cat > /etc/init/cgroup-lite.conf <<-'EOF'
description "mount available cgroup filesystems"
author "Serge Hallyn <serge.hallyn@canonical.com>"

start on mounted MOUNTPOINT=/sys/fs/cgroup

pre-start script
	test -x /bin/cgroups-mount || { stop; exit 0; }
	test -d /sys/fs/cgroup || { stop; exit 0; }
	/bin/cgroups-mount
	/usr/sbin/cgconfigparser -l /etc/cgconfig.conf
end script

post-stop script
	if [ -x /bin/cgroups-umount ]
	then
		/bin/cgroups-umount
	fi
end script
EOF

echo "Creating /etc/init/cgrulesengd.conf"
cat > /etc/init/cgrulesengd.conf <<-'EOF'
# cgrulesengd

description "cgrulesengd"
author "Serge Hallyn &lt;serge.hallyn@canonical.com&gt;"

start on started cgroup-lite
stop on stopped cgroup-lite

pre-start script
test -x /usr/sbin/cgrulesengd || { stop; exit 0; }
end script

script
# get default options
OPTIONS=""
CGRED_CONF=/etc/cgrules.conf
if [ -r "/etc/default/cgrulesengd" ]; then
. /etc/default/cgrulesengd
fi

# Don't run if no configuration file
if [ ! -s "$CGRED_CONF" ]; then
echo "Cgred unconfigured"
stop
exit 0
fi

# Make sure the kernel supports cgroups
# This check is retained from the original sysvinit job, but should
# be superfluous since we depend on cgconfig running, which will
# have mounted this.
grep -q "^cgroup" /proc/mounts || { stop; exit 0; }

exec /usr/sbin/cgrulesengd --nodaemon $OPTIONS
end script
EOF

echo "------------------------------------------------------"

service cgroup-lite restart
service cgrulesengd restart
                                                  
} # this ensures the entire script is downloaded #