#!/bin/bash
################################################################################
# echo wrappers
INFO(){ echo -n "INFO: "; echo "$@" ;}
WARN(){ echo -n "WARN: "; echo "$@" ;}
ERRO(){ echo -n "ERRO: "; echo "$@" ; exit 1;}

################################################################################
# Helpers
YN(){
    case "$1" in
        Yes|Y|1|true) return 0 ;;
        *) return 1 ;;
    esac
}

write(){
    DATA="$1" FILE="$2"
    [ -z "$DATA" ] && return
    [ -z "$FILE" ] && return
    echo "$DATA" > "$FILE"
}

help(){
    echo "$0 start|stop|status"
    echo "   start  - init daemon"
    echo "   stop   - stop daemon"
    echo "   status - show some swap status info"
}

get_fs_type(){ stat -f "$1" | grep Type: | tr -s ' ' | cut -d' ' -f7; }

################################################################################
# Systemd swap unit generator
systemd_swapon(){
        UNIT_PATH=/run/systemd/system/
        export What Priority Options
        for i in "$@"; do
            case $i in
                What=*)     What="${i//What=/}" ;;
                Priority=*) Priority="${i//Priority=/}" ;;
                Options=*)  Options="${i//Options=/}" ;;
            esac
        done
        [ ! -z "$What" ] || return 1
        What="$(realpath $What)"
        UNIT_NAME="$(systemd-escape -p $What)"
        {
                echo '[Swap]'
                echo "What=$What"
                [ -z "$Priority" ] || echo "Priority=$Priority"
                [ -z "$Options"  ] || echo "Options=$Options"
        } > "$UNIT_PATH/$UNIT_NAME.swap"
        systemctl start $UNIT_NAME.swap
}

################################################################################
# Initialization
check_root_rights(){ [ "$UID" == "0" ] || ERRO "Script must be run as root!"; }

# get cpu count from cpuinfo
cpu_count=$(nproc)
# get total ram size for meminfo
ram_size=$(awk '/MemTotal:/ { print $2 }' /proc/meminfo)
# get page size
PAGE_SIZE=$(getconf PAGESIZE)

CONF=/etc/systemd/swap.conf
WORK_DIR=/run/systemd/swap
B_CONF=$WORK_DIR/swap.conf

case "$1" in
    start)
        INFO "Check config"
        [ ! -f $CONF ] && ERRO "Missing config: /etc/systemd/swap.conf - try to reinstall package"
        check_root_rights
        mkdir -p $WORK_DIR
        [ -f $B_CONF ] && ERRO "systemd-swap already started!"
        INFO "Backup config"
        cp $CONF $B_CONF
        if [ -d /etc/systemd/swap.conf.d/ ]; then
                for file in /etc/systemd/swap.conf.d/*.conf; do
                        [ ! -f "$file" ] && continue
                        INFO "Load: $file"
                        cat "$file" >> $B_CONF
                done
        fi
        INFO "Load config"
        # shellcheck source=/run/systemd/swap/swap.conf
        . $B_CONF

        zswap_enabled=${zswap_enabled:-0}
        if YN $zswap_enabled; then
            [ ! -d /sys/module/zswap ] && ERRO "Zswap - not supported on current kernel"
            ZSWAP_P=/sys/module/zswap/parameters/
            INFO "Zswap: backup current configuration: start"
            mkdir -p $WORK_DIR/zswap/
            for file in $ZSWAP_P/*; do
                cp "$file" "$WORK_DIR/zswap/$(basename $file)"
            done
            INFO "Zswap: backup current configuration: complete"
            INFO "Zswap: set new parameters: start"
            write ${zswap_enabled:-1}           $ZSWAP_P/enabled
            write ${zswap_compressor:-lz4}      $ZSWAP_P/compressor
            write ${zswap_max_pool_percent:-25} $ZSWAP_P/max_pool_percent
            write ${zswap_zpool:-zbud}          $ZSWAP_P/zpool
            INFO "Zswap: set new parameters: complete"
        fi

        zram_enabled=${zram_enabled:-0}
        if YN $zram_enabled; then
            [ -z "$zram_size" ] && zram_size=$((ram_size/4))K
            zram_streams=${zram_streams:-$cpu_count}
            zram_alg=${zram_alg:-"lz4"}
            zram_prio=${zram_prio:-"32767"}
            zram_dev=""
            INFO "Zram: check availability"
            if [ ! -d /sys/module/zram ]; then
                INFO "Zram: not part of kernel, trying to find zram module"
                modprobe -n zram || ERRO "Zram: can't find zram module!"
                # Wrapper, for handling zram initialization problems
                for (( i = 0; i < 10; i++ )); do
                    if [ ! -d /sys/module/zram ]; then
                        modprobe zram
                        sleep 1
                    fi
                done
                INFO "Zram: module successfully loaded"
            else
                INFO "Zram: module already loaded"
            fi
            for (( i = 0; i < 10; i++ )); do
                INFO "Zram: trying to initialize free device"
                # zramctl is a external program -> return name of first free device
                TMP=$(mktemp)
                zramctl -f -a $zram_alg -t $zram_streams -s $zram_size &> $TMP
                read -r OUTPUT < $TMP
                rm $TMP
                case "$OUTPUT" in
                    *"failed to reset: Device or resource busy"*) sleep 1 ;;
                    *"zramctl: no free zram device found"*)
                        WARN "Zram: zramctl can't find free device"
                        INFO "Zram: using workaround hook for hot add"
                        [ ! -f /sys/class/zram-control/hot_add ] && \
                            ERRO "Zram: this kernel does not support hot add zram device, please use 4.2+ kernels or see modinfo zram and make a modprobe rule"
                        NEW_ZRAM=$(cat /sys/class/zram-control/hot_add)
                        INFO "Zram: success: new device /dev/zram$NEW_ZRAM"
                    ;;
                    /dev/zram*)
                        [ -b "$OUTPUT" ] || continue
                        zram_dev="$OUTPUT"
                        break
                    ;;
                esac
            done
            INFO "Zram: initialized: $zram_dev"
            mkdir -p $WORK_DIR/zram/
            mkswap "$zram_dev" &> /dev/null && {
                    systemd_swapon What="$zram_dev" Options=discard Priority=$zram_prio && \
                    ln -s $zram_dev $WORK_DIR/zram/
            }
        fi

        swapfu_enabled=${swapfu_enabled:-0}
        if YN $swapfu_enabled; then
            swapfu_size=${swapfu_size:-"${ram_size}K"}
            swapfu_path=${swapfu_path:-"/var/lib/systemd-swap/swapfu"}
            swapfu_prio=${swapfu_prio:-"-1024"}

            [ -d "$swapfu_path" ] && swapfu_path="${swapfu_path}/${RANDOM}_swapfu"

            INFO "swapFU: create empty file"
            mkdir -p "$(dirname $swapfu_path)"

            swapfu_nocow=${swapfu_nocow:-1}
            if YN $swapfu_nocow; then
                INFO "swapFU: mark file $swapfu_path nocow"
                touch $swapfu_path
                chattr +C $swapfu_path
            fi

            INFO "swapFU: max file size: 0 -> $swapfu_size"
            swapfu_preallocated=${swapfu_preallocated:-0}
            if YN $swapfu_preallocated; then
                INFO "swapFU: preallocate $swapfu_size"
                fallocate -l $swapfu_size $swapfu_path
            else
                truncate -s $swapfu_size $swapfu_path
            fi

            INFO "swapFU: attach $swapfu_path to free loop"
            swapfu_loop=$(losetup -f --show $swapfu_path)
            INFO "swapFU: using $swapfu_loop"
            # loop uses file descriptor, the file still exists,
            # but does not have a path like O_TMPFILE
            # When loop detaches a file, the file will be deleted.
            rm $swapfu_path

            INFO "swapFU: initialize swap device $swapfu_loop"
            mkswap $swapfu_loop &> /dev/null

            if YN $swapfu_preallocated; then
                INFO "swapFU: swapon - discard: disabled"
                systemd_swapon What="$swapfu_loop" Priority=$swapfu_prio
            else
                INFO "swapFU: swapon - discard: enabled"
                systemd_swapon What="$swapfu_loop" Options=discard Priority=$swapfu_prio
            fi

            mkdir -p $WORK_DIR/swapfu
            ln -s $swapfu_loop $WORK_DIR/swapfu/

            # set autoclear flag
            losetup -d $swapfu_loop

            swapfu_directio=${swapfu_directio:-0}
            if YN $swapfu_directio; then
                INFO "swapFU: enable directio for $swapfu_loop"
                losetup --direct-io=on  $swapfu_loop
            else
                INFO "swapFU: disable directio for $swapfu_loop"
                losetup --direct-io=off $swapfu_loop
            fi
        fi

        swapfc_enabled=${swapfc_enabled:-0}
        if YN $swapfc_enabled; then
            get_free_swap_perc(){
                    # +1 prevent devide by zero
                    total="$(grep SwapTotal: /proc/meminfo | grep -o -e '[0-9]*')"
                    free="$(grep SwapFree: /proc/meminfo | grep -o -e '[0-9]*')"
                    total=$((total*1024))
                    free=$((free*1024))
                    echo $(( (${free}*100)/(${total}+1) ));
            }
            mkswap_w(){
                chunk_size="$1" file="$2"
                fallocate -l "$chunk_size" "$file"
                # File must not contain holes for XFS/f2fs & etc
                case "$(get_fs_type $file)" in
                    ext2|ext3|ext4);;
                    btrfs) ERRO "swapFC: Btrfs doesn't support swap files" ;;
                    *) shred -n1 -z "$file" ;;
                esac
                chmod 0600 "$file"
                mkswap "$file" &> /dev/null
            }
            chunk_size=${swapfc_chunk_size:-"256M"}
            max_count=${swapfc_max_count:-"16"}
            free_swap_perc=${swapfc_free_swap_perc:-"15"}
            swapfc_path=${swapfc_path:-"/var/lib/systemd-swap/swapfc/"}
            swapfc_frequency=${swapfc_frequency:-"1s"}
            swapfc_max_count=${swapfc_max_count:-"16"}

            if (( "$swapfc_max_count" > 32 )) || (( 1 > "$swapfc_max_count" )); then
                WARN "swapfc_max_count must be in range 1..32, reset to 1"
                swapfc_max_count=1
            fi

            mkdir -p "$swapfc_path"
            mkdir -p "$WORK_DIR/swapfc/"
            touch    "$WORK_DIR/swapfc/.lock"

            {
                allocated=0
                while sleep $swapfc_frequency; do
                    [ -f $WORK_DIR/swapfc/.lock ] || break
                    curr_free_swap_perc=$(get_free_swap_perc)
                    if (( $curr_free_swap_perc < $free_swap_perc )) && (( allocated < max_count )); then
                        allocated=$((allocated+1))
                        INFO "swapFC: free swap: $curr_free_swap_perc < $free_swap_perc - allocate chunk: $allocated"
                        mkswap_w "$chunk_size" "$swapfc_path/$allocated"
                        systemd_swapon What="$swapfc_path/$allocated" && \
                            ln -s "$swapfc_path/$allocated" "$WORK_DIR/swapfc/$allocated"
                    fi
                    if (( $(get_free_swap_perc) > $((free_swap_perc+40)) )) && (( allocated > 2 )); then
                        INFO "swapFC: free swap: $curr_free_swap_perc > $((free_swap_perc+40)) - freeup chunk: $allocated"
                        swapoff "$WORK_DIR/swapfc/$allocated" && \
                            rm "$swapfc_path/$allocated" && \
                            rm "$WORK_DIR/swapfc/$allocated"
                        allocated=$((allocated-1))
                    fi
                done
            } &
        fi

        swapd_auto_swapon=${swapd_auto_swapon:-1}
        if YN $swapd_auto_swapon; then
            swapd_prio=${swapd_prio:-1024}
            INFO "swapD: searching swap devices"
            mkdir -p $WORK_DIR/swapd/
            for device in $(blkid -t TYPE=swap -o device | grep -vE '(zram|loop)'); do
                for used_device in $(swapon --show=NAME --noheadings); do
                    [ "$device" == "$used_device" ] && unset device
                done
                [ ! -b "$device" ] && continue
                systemd_swapon What="$device" Options=discard Priority="$swapd_prio" || continue
                ln -s $device $WORK_DIR/swapd/ || continue
                INFO "swapD: enabled device: $device"
                swapd_prio=$((swapd_prio-1))
            done
        fi
    ;;
    stop)
        check_root_rights
        [ ! -f $B_CONF ] && ERRO "systemd-swap failed to start!"
        INFO "Load config"
        # shellcheck source=/run/systemd/swap/swap.conf
        . $B_CONF

        if [ -d $WORK_DIR/swapd/ ]; then
            for device in $WORK_DIR/swapd/*; do
                [ ! -b "$device" ] && continue
                INFO "swapD: swapoff: $device start"
                swapoff $device || continue
                rm $device      || continue
                INFO "swapD: swapoff: $device complete"
            done
        fi

        if [ -d $WORK_DIR/swapfc/ ]; then
            swapfc_path=${swapfc_path:-"/var/lib/systemd-swap/swapfc/"}
            rm -f $WORK_DIR/swapfc/.lock
            # wait two times
            sleep $swapfc_frequency
            sleep $swapfc_frequency
            for file in $WORK_DIR/swapfc/*; do
                [ -f "$file" ] || continue
                swapoff $file  || continue
                rm -f $file "$swapfc_path/$(basename $file)"
            done
            [ -d $swapfc_path ] && rmdir $swapfc_path
        fi

        if [ -d $WORK_DIR/swapfu/ ]; then
            for device in $WORK_DIR/swapfu/*; do
                [ -b "$device" ] || continue
                swapoff $device  || continue
                rm $device       || continue
                INFO "swapF: stopped /dev/$(basename $device)"
            done
        fi

        if [ -d $WORK_DIR/zram/ ]; then
            for zram in $WORK_DIR/zram/*; do
                [ -b $zram ]    || continue
                INFO "Zram: removing: /dev/$(basename $zram)"
                swapoff $zram   || continue
                zramctl -r "$(basename $zram)" || continue
                rm $zram        || continue
                INFO "Zram: removed: /dev/$(basename $zram)"
            done
        fi

        if [ -d $WORK_DIR/zswap/ ]; then
            ZSWAP_P=/sys/module/zswap/parameters/
            INFO "Zswap: restore configuration: start"
            for file in $WORK_DIR/zswap/*; do
                cp "$file" "$ZSWAP_P/$(basename $file)"
            done
            INFO "Zswap: restore configuration: complete"
        fi

        rm -rf $WORK_DIR
    ;;
    status)
        check_root_rights
        if [ -d /sys/module/zswap ]; then
                echo Zswap:
                read USED_BYTES   < /sys/kernel/debug/zswap/pool_total_size
                USED_PAGES=$((USED_BYTES/$PAGE_SIZE))
                read STORED_PAGES < /sys/kernel/debug/zswap/stored_pages
                STORED_BYTES=$((STORED_PAGES*$PAGE_SIZE))
                SWAP_STORE_USED_BYTES=( $(swapon --bytes | awk '{print $4}' | grep -v USED) )
                RATIO=0
                ((STORED_PAGES>0)) && RATIO=$((USED_PAGES*100/STORED_PAGES))
                SUM=0
                if ((${#SWAP_STORE_USED_BYTES[@]} > 0)); then
                        for val in "${SWAP_STORE_USED_BYTES[@]}"; do
                                SUM=$((SUM+val))
                        done
                fi
                {
                        grep . /sys/module/zswap/parameters/* | \
                                cut -d'/' -f 6 | \
                                sed -e 's/:/ /g' -e 's/^/. /g'
                } | column -t
                {
                        grep . /sys/kernel/debug/zswap/* | \
                                cut -d'/' -f6 | \
                                sed -e 's/:/ /g' -e 's/^/. . /g' | column -t
                        echo . . compress_ratio ${RATIO}%
                        ((SUM > 0)) && \
                                echo . . zswap_store/swap_store ${STORED_BYTES}/${SUM} $((STORED_BYTES*100/SUM))%
                } | column -t

        else
                WARN "Zswap unavailable"
        fi
        if zramctl | grep -q '\[SWAP\]'; then
                echo ZRam:
                zramctl  | \
                        grep -e '^NAME\|\[SWAP\]' | \
                        sed 's/^/. /g' | \
                        column -t \
                                --table-columns \.,NAME,ALGORITHM,DISKSIZE,DATA,COMPR,TOTAL,STREAMS,MOUNTPOINT \
                                -H MOUNTPOINT | \
                                uniq
        fi
        if [ -d /run/systemd/swap/swapfu ]; then
                {
                        echo SwapFU:
                        swapon | grep NAME | sed 's/^/. /g'
                        for loop in /run/systemd/swap/swapfu/*; do
                                [ -b $loop ] || continue
                                swapon | grep "$(basename $loop)"
                        done | sed 's/^/. /g'
                } | column -t
        fi
        if [ -d /run/systemd/swap/swapfc ]; then
                {
                        echo SwapFC:
                        swapon --show TYPE | grep -e 'NAME\|file' | sed 's/^/. /g'
                } | column -t
        fi
        if [ -d /run/systemd/swap/swapd ]; then
                {
                        echo SwapD:
                        swapon | grep NAME | sed 's/^/. /g'
                        for dev in /run/systemd/swap/swapd/*; do
                                [ -b $dev ] || continue
                                swapon | grep "$(basename $dev)"
                        done | sed 's/^/. /g'
                } | column -t
        fi
    ;;
    *) help ;;
esac
