#!/bin/bash
# This script does the following:
# zramswap start:
#  * Space is assigned to each zram device, then swap is initialized on
#    there
# zramswap stop:
#  * Undo start
#  * Also attempts to remove zram module at the end
# TODO:
# * Migrate to using zramctl from util-linux for the setup,
#   (this will close debian bug #917643):
#   then also:
#   - add option for compression algorythm
#   - ammount of compression streams
#   - Make use of the zramctl stats too

function start {
    #Set some defaults:
    ALLOCATION=10240 # ZRAM Swap you want assigned, in MiB
    PRIORITY=100   # Swap priority, see swapon(2) for more details

    # Get amount of available CPU cores, set to 1 if not detected correctly
    if [ ! -f /proc/cpuinfo ]; then
        echo "WARNING: Can't find /proc/cpuinfo, is proc mounted?"
        echo "         Using a single core for zramswap..."
        CORES=1
    else
        CORES=$(grep -c ^processor /proc/cpuinfo)
    fi

    # Override  above from config file, if it exists
    if [ -f /etc/default/zramswap ]; then
        . /etc/default/zramswap
    fi

    ALLOCATION=$((ALLOCATION * 1024 * 1024)) # convert amount from MiB to bytes

    if [ -n "$PERCENTAGE" ]; then
        totalmemory=$(awk '/MemTotal/{print $2}' /proc/meminfo) # in KiB
        ALLOCATION=$((totalmemory * 1024 * $PERCENTAGE / 100))
    fi

    # Initialize zram devices
    modprobe zram num_devices=$CORES

    # Assign memory to zram devices, initialize swap and activate
    # Decrementing $CORE, because cores start counting at 0
    for CORE in $(seq 0 $(($CORES - 1))); do
	echo "zstd" > /sys/block/zram$CORE/comp_algorithm
        echo $(($ALLOCATION / $CORES)) > /sys/block/zram$CORE/disksize
        mkswap /dev/zram$CORE
        swapon -p $PRIORITY /dev/zram$CORE
    done
}

function status {
    orig_data_size="0"
    for file in /sys/block/zram*/*_data_size ; do
        if [ $file = "/sys/block/zram*/*_data_size" ]; then
            compress_ratio="0"
            break
        fi
        read file_content < $file
        what=$(basename $file)
        eval "$what=\$(($what + $file_content))"
        compress_ratio=$(echo "scale=2; $orig_data_size / $compr_data_size" | bc)
    done
    echo "compr_data_size: $((compr_data_size / 1024)) KiB"
    echo "orig_data_size:  $((orig_data_size  / 1024)) KiB"
    echo "compression-ratio: $compress_ratio"
}

function stop {
    for swapspace in $(swapon -s | awk '/zram/{print $1}'); do
        swapoff $swapspace
    done
    modprobe -r zram
}

function usage {
    echo "Usage:"
    echo "   zramswap start - start zram swap"
    echo "   zramswap stop - stop zram swap"
    echo "   zramswap status - prints some statistics"
}

if [ "$1" = "start" ]; then
    start
fi

if [ "$1" = "stop" ]; then
    stop
fi

if [ "$1" = "status" ]; then
    status
fi

if [ "$1" = "" ]; then
    usage
fi