#!/bin/bash

# Script for parsing cgroup information
#
# This script will read some limits from the cgroup system and parse
# them, printing out "VARIABLE=VALUE" on each line for every limit that is
# successfully read. Output of this script can be directly fed into
# bash's export command. Recommended usage from a bash script:
#
#     set -o errexit
#     export_vars=$(cgroup-limits) ; export $export_vars
#
# Variables currently supported:
#     MAX_MEMORY_LIMIT_IN_BYTES
#         Maximum possible limit MEMORY_LIMIT_IN_BYTES can have. This is
#         currently a constant value of 9223372036854775807.
#     MEMORY_LIMIT_IN_BYTES
#         Maximum amount of user memory in bytes. If this value is set
#         to the same value as MAX_MEMORY_LIMIT_IN_BYTES, it means that
#         there is no limit set. The value is taken from
#         /sys/fs/cgroup/memory/memory.limit_in_bytes
#     RAM_SIZE_IN_BYTES
#         Available memory in bytes. The value is taken from /proc/meminfo.
#     NUMBER_OF_CORES
#         Number of detected CPU cores that can be used. This value is
#         calculated from /sys/fs/cgroup/cpuset/cpuset.cpus
#     NO_MEMORY_LIMIT
#         Set to "true" if MEMORY_LIMIT_IN_BYTES is so high that the caller
#         can act as if no memory limit was set. Undefined otherwise.

get_memory_limit() {
    local limit="$(cat /sys/fs/cgroup/memory/memory.limit_in_bytes)"
    if [ $? -eq 0 ]; then
        echo "$limit"
    else
        echo "Warning: Can't detect memory limit from cgroups" >&2
    fi
}

get_meminfo_ram_size() {
    local limit="$(cat /proc/meminfo | while read key val unit rest; do if [ "$key" = "MemTotal:" ]; then echo "$val $unit"; break; fi; done)"
    if [ $? -eq 0 ]; then
        local val="${limit/ *}"
        local unit="${limit/* }"
        case "$unit" in
            kB)
                val=$(($val * 1024))
                ;;
            MB)
                val=$(($val * 1024 * 1024))
                ;;
            GB)
                val=$(($val * 1024 * 1024 * 1024))
                ;;
        esac
        echo "$val"
    else
        echo "Warning: Can't detect RAM site from /proc/meminfo." >&2
    fi
}

get_number_of_cores() {
    local core_count=0
    local group
    local sysfs_cpus="$(cat /sys/fs/cgroup/cpuset/cpuset.cpus)"

    if [ $? -eq 0 ]; then
        for group in ${sysfs_cpus//,/ }; do
            local from="${group%-*}"
            local to="${group#*-}"
            core_count=$(($core_count + $to - $from + 1))
        done
        echo "$core_count"
    else
        echo "Warning: Can't detect number of CPU cores from cgroups" >&2
    fi
}

MAX_MEMORY_LIMIT_IN_BYTES=9223372036854775807
MEMORY_LIMIT_IN_BYTES="$(get_memory_limit)"
RAM_SIZE_IN_BYTES="$(get_meminfo_ram_size)"
NUMBER_OF_CORES="$(get_number_of_cores)"

echo "MAX_MEMORY_LIMIT_IN_BYTES=$MAX_MEMORY_LIMIT_IN_BYTES"

if [ -n "$MEMORY_LIMIT_IN_BYTES" ]; then
    echo "MEMORY_LIMIT_IN_BYTES=$MEMORY_LIMIT_IN_BYTES"
    if [ "$MEMORY_LIMIT_IN_BYTES" -ge "$MAX_MEMORY_LIMIT_IN_BYTES" ]; then
        echo "NO_MEMORY_LIMIT=true"
    fi
fi

if [ -n "$RAM_SIZE_IN_BYTES" ]; then
    echo "RAM_SIZE_IN_BYTES=$RAM_SIZE_IN_BYTES"
fi

if [ -n "$NUMBER_OF_CORES" ]; then
    echo "NUMBER_OF_CORES=$NUMBER_OF_CORES"
fi