mm_limit

#!/bin/bash

if mount | grep -qw cgroup2; then
    ver=2
elif mount | grep -qw cgroup; then
    ver=1
else
    echo "No cgroup detected."
    exit 1
fi

cg_name="mm_limit_cgroup"

if [[ $ver == 1 ]]; then
    cg_root="/sys/fs/cgroup/memory"
else
    cg_root="/sys/fs/cgroup"
fi

cg_path="$cg_root/$cg_name"

if [[ $ver == 1 ]]; then
    mem_high_path="$cg_path/memory.limit_in_bytes"
else
    mem_high_path="$cg_path/memory.high"
fi

proc_path="$cg_path/cgroup.procs"

# By default, allow 4MB mem usage
mem_high_default=4M

pid="$1"
mem_high="$2"

err()
{
    echo "$@"
    exit 1
}

usage()
{
    cat <<EOF

usage: $0 <pid> [mem_high]

Create a shared cgroup "$cg_name" and put process into the group.  When
mem_high is specified, limit memory usage of the group to 'mem_high'.  By
default, mem_high=$mem_high_default.

The user must be in the sudoers list.

EOF
    exit 1
}

if [[ -z "$mem_high" ]]; then
    mem_high=$mem_high_default
fi

if [[ -z "$pid" ]]; then
    usage
fi

if [[ ! -d "$cg_path" ]]; then
    echo "Creating directory $cg_path"
    sudo mkdir $cg_path || err "Failed to create $cg_path"
fi

if ! [[ -f "$mem_high_path" ]]; then
    err "Didn't find path $mem_high_path"
fi

echo "Setting $mem_high_path to $mem_high..."
echo $mem_high | sudo tee $mem_high_path &> /dev/null || \
    err "Failed to setup memory.high"

echo "Moving PID $pid into cgroup..."
echo $pid | sudo tee $proc_path &> /dev/null ||
    err "Failed to add PID $pid to cgroup"

echo "All done"