#!/bin/bash
################################################################################
# echo wrappers
INFO(){ echo "INFO: $*";}
WARN(){ echo "WARN: $*";}
ERRO(){ echo "ERRO: $*"; exit 1;}
INFO_E(){ echo "INFO: $*"; exit 0;}

################################################################################
# Helpers
YN(){
        l="${1,,}"
        [[ "$l" =~ (yes|y|1|true) ]] && return 0
        return 1
}

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

# Files with one byte data will use 1 page of ram in tmpfs,
# i.e. 4KiB on x86_64
# that a madness, store short string in symbol links
write_l(){
        { [ -z "$1" ] || [ -z "$2" ]; } && return
        DATA="$1" FILE="$2"
        DATA=$(echo "$DATA" | base64)
        ln -sfr "/$DATA" "$FILE"
}

read_l(){
        [ -z "$1" ] && return
        FILE="$1"
        DATA=$(basename "$(realpath $FILE)" )
        echo $DATA | base64 -d
}

get_fs_type(){ stat -fc "%T" "$1"; }
AMI_ROOT(){ [ "$UID" == "0" ] || ERRO "Script must be run as root!"; }
FIND_SWAP_UNITS(){ find /run/systemd/ -type f -name "*.swap"; }
GET_WHAT_FROM_SWAP_UNIT(){ grep What $1 | cut -d'=' -f2; }
help(){
    echo "$0 start|stop|status"
    echo "   start  - init daemon"
    echo "   stop   - stop daemon"
    echo "   status - show some swap status info"
}

################################################################################ы
# Global vars
readonly RUN_SYSD=/run/systemd/
readonly NCPU=$(nproc)
readonly RAM_SIZE=$(($(grep MemTotal: /proc/meminfo | grep -o '[0-9]*')*1024))
readonly PAGE_SIZE=$(getconf PAGESIZE)
readonly ETC_SYSD=/etc/systemd/
readonly CONFIG=$ETC_SYSD/swap.conf
readonly WORK_DIR=/run/systemd/swap

readonly LOCK_STARTED="$WORK_DIR/.started"
readonly ZSWAP_M="/sys/module/zswap"
readonly ZSWAP_M_P="/sys/module/zswap/parameters"

################################################################################
# swap unit generator
gen_swap_unit(){
        export What Priority Options Tag Type
        for i in "$@"; do
            case $i in
                What=*)     What="${i//What=/}" ;;
                Priority=*) Priority="${i//Priority=/}" ;;
                Options=*)  Options="${i//Options=/}"   ;;
                Tag=*)      Tag="${i//Tag=}"          ;;
            esac
        done
        [ ! -z "$What" ] || return 1

        What="$(realpath $What)"
        Type="File" # Assume that a file by default
        if [ -b "$What" ]; then
                Type="Block/Partition"
                [[ "$What" =~ loop ]] && Type="File"
        fi

        UNIT_NAME="$(systemd-escape -p $What)"
        UNIT_PATH="$RUN_SYSD/system/$UNIT_NAME.swap"
        {
                echo '[Unit]'
                echo "Description=Swap $Type"
                echo 'Documentation=https://github.com/Nefelim4ag/systemd-swap'
                echo
                echo "# Generated by systemd-swap"
                echo "# Tag=$Tag"
                echo
                echo '[Swap]'
                echo "What=$What"
                [ -z "$Priority" ] || echo "Priority=$Priority"
                [ -z "$Options"  ] || echo "Options=$Options"
        } > "$UNIT_PATH"

        ln -srf "$UNIT_PATH" $RUN_SYSD/system/swap.target.wants/

        if [[ "$Type" == "File" ]]; then
                ln -srf "$UNIT_PATH" $RUN_SYSD/system/local-fs.target.wants/
        fi

        echo "${UNIT_NAME}.swap"
}

################################################################################
# INIT
################################################################################

AMI_ROOT

mkdir -p $WORK_DIR \
         $RUN_SYSD/system/local-fs.target.wants \
         $RUN_SYSD/system/swap.target.wants

case "$1" in
        start)
                [ -f $LOCK_STARTED ] && ERRO "$0 already started"
                touch $LOCK_STARTED

                INFO "Load: $CONFIG"
                # shellcheck source=/run/systemd/swap/swap.conf
                . $CONFIG || ERRO "Problems while load of $CONFIG"

                if [ -d $ETC_SYSD/swap.conf.d ]; then
                        for L_CONF in  $ETC_SYSD/swap.conf.d/*; do
                                [ ! -f "$L_CONF" ] && continue
                                # shellcheck source=/run/systemd/swap/swap.conf
                                . $L_CONF || ERRO "Problems while load of $L_CONF"
                        done
                fi

                zswap_enabled=${zswap_enabled:-0}
                if YN $zswap_enabled; then
                        [ ! -d $ZSWAP_M ] && ERRO "Zswap - not supported on current kernel"

                        zswap_compressor=${zswap_compressor:-lz4}
                        zswap_max_pool_percent=${zswap_max_pool_percent:-25}
                        zswap_zpool=${zswap_zpool:-zbud}

                        INFO "Zswap: backup current configuration: start"
                        mkdir -p $WORK_DIR/zswap/
                        for file in $ZSWAP_M_P/*; do
                                read -r VAL < $file
                                write_l "$VAL" "$WORK_DIR/zswap/$(basename $file)"
                        done
                        INFO "Zswap: backup current configuration: complete"
                        INFO "Zswap: set new parameters: start"
                        INFO "Zswap: Enable: $zswap_enabled, Comp: $zswap_compressor,  Max pool %: $zswap_max_pool_percent, Zpool: $zswap_zpool"
                        write ${zswap_enabled}          $ZSWAP_M_P/enabled
                        write ${zswap_compressor}       $ZSWAP_M_P/compressor
                        write ${zswap_max_pool_percent} $ZSWAP_M_P/max_pool_percent
                        write ${zswap_zpool}            $ZSWAP_M_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))
                        zram_streams=${zram_streams:-$NCPU}
                        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
                                        [ -d /sys/module/zram ] && break
                                        modprobe zram
                                        sleep 1
                                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 -f $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"
                                                read -r NEW_ZRAM < /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"
                        mkswap "$zram_dev" &> /dev/null && \
                        UNIT_NAME=$(gen_swap_unit What="$zram_dev" Options=discard Priority=$zram_prio Tag=zram)
                        systemctl daemon-reload
                        systemctl start $UNIT_NAME
                fi

                swapfc_enabled=${swapfc_enabled:-0}
                if YN $swapfc_enabled; then
                        chunk_size=${swapfc_chunk_size:-"512M"}
                        swapfc_max_count=${swapfc_max_count:-"8"}
                        swapfc_path=${swapfc_path:-"/var/lib/systemd-swap/swapfc/"}
                        swapfc_frequency=${swapfc_frequency:-"1s"}
                        swapfc_force_use_loop=${swapfc_force_use_loop:-"0"}
                        swapfc_nocow=${swapfc_nocow:-"1"}
                        swapfc_directio=${swapfc_directio:-"1"}
                        swapfc_force_preallocated=${swapfc_force_preallocated:-"0"}
                        swapfc_free_swap_perc=${swapfc_free_swap_perc:-"15"}
                        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) ));
                        }

                        to_bytes(){ numfmt --to=none --from=iec "$1"; }

                        chunk_size=$(to_bytes $chunk_size)
                        BLOCK_SIZE=$(stat -f -c %s "$swapfc_path")
                        FSTYPE="$(get_fs_type $swapfc_path)"

                        YN $swapfc_force_use_loop && FSTYPE=btrfs

                        [ "$FSTYPE" != btrfs ] && swapfc_force_preallocated=1

                        check_ENOSPC(){
                                path="$1"
                                # Check free space
                                # For avoid problems on swap io + ENOSPC
                                FREE_BLOCKS=$(stat -f -c %f "$path" )
                                FREE_BYTES=$((FREE_BLOCKS*BLOCK_SIZE))
                                # also try leave some free space
                                FREE_BYTES=$((FREE_BYTES - chunk_size))
                                ((FREE_BYTES < chunk_size)) && return 0
                                return 1
                        }

                        prepare_swapfile(){
                                chunk_size="$1" file="$2"
                                touch "$file"
                                chmod 0600 "$file"

                                YN $swapfc_nocow && chattr +C "$file"

                                if YN $swapfc_force_preallocated; then
                                        fallocate -l "$chunk_size" "$file"
                                else
                                        truncate -s "$chunk_size" "$file"
                                fi

                                losetup_w(){
                                        file="$1"
                                        if YN $swapfc_directio; then
                                                losetup -f --show --direct-io=on  "$file"
                                        else
                                                losetup -f --show --direct-io=off "$file"
                                        fi
                                        # 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 "$file"
                                }

                                RET="$file"
                                case "$FSTYPE" in
                                        ext2|ext3|ext4) ;;
                                        btrfs) RET=$(losetup_w $file) ;;
                                        *)
                                                # -z rewrite file second time
                                                # By default shred gen random
                                                # Speedup rewrite by rewrite
                                                # one time with zeroes
                                                shred -n1 --random-source=/dev/zero "$file"
                                        ;;
                                esac
                                echo "$RET"
                        }

                        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" \
                                 "$WORK_DIR/swapfc/"
                        touch    "$WORK_DIR/swapfc/.lock"
                        {
                                allocated=0
                                free_threshold_up=$swapfc_free_swap_perc
                                free_threshold_down=$((swapfc_free_swap_perc+40))
                                while [ -f $WORK_DIR/swapfc/.lock ] && sleep $swapfc_frequency; do
                                        curr_free_swap_perc=$(get_free_swap_perc)
                                        if (( curr_free_swap_perc < free_threshold_up )) && (( allocated < swapfc_max_count )); then
                                                if check_ENOSPC "$swapfc_path"; then
                                                        WARN "swapFC: ENOSPC"
                                                        continue
                                                fi
                                                allocated=$((allocated+1))
                                                INFO "swapFC: free swap: $curr_free_swap_perc < $free_threshold_up - allocate chunk: $allocated"
                                                swapfile=$(prepare_swapfile "$chunk_size" "$swapfc_path/$allocated")
                                                mkswap -L SWAP_${FSTYPE}_${allocated} "$swapfile" | sed ':a;N;$!ba;s/)\n/), /g'
                                                if YN $swapfc_force_preallocated; then
                                                        UNIT_NAME=$(gen_swap_unit What="$swapfile" Tag="swapfc_$allocated")
                                                else
                                                        UNIT_NAME=$(gen_swap_unit What="$swapfile" Options=discard Tag="swapfc_$allocated")
                                                fi
                                                systemctl daemon-reload
                                                systemctl start $UNIT_NAME
                                                if [ -b "$swapfile" ]; then
                                                        losetup -d "$swapfile"
                                                fi
                                        fi
                                        (( allocated <= 2 )) && continue
                                        if (( curr_free_swap_perc > free_threshold_down )); then
                                                INFO "swapFC: free swap: $curr_free_swap_perc > $free_threshold_down - freeup chunk: $allocated"
                                                FIND_SWAP_UNITS | while read -r UNIT_PATH; do
                                                        if grep -q swapfc_$allocated $UNIT_PATH; then
                                                                DEV=$(GET_WHAT_FROM_SWAP_UNIT $UNIT_PATH)
                                                                UNIT_NAME=$(basename $UNIT_PATH)
                                                                systemctl stop "$UNIT_NAME" || swapoff "$DEV"
                                                                rm -f "$UNIT_NAME"
                                                                [ -f $DEV ] && rm -f $DEV
                                                        fi
                                                done
                                                allocated=$((allocated-1))
                                        fi
                                done
                        } &
                fi

                swapd_auto_swapon=${swapd_auto_swapon:-1}
                swapd_prio=${swapd_prio:-1024}
                if YN $swapd_auto_swapon; then
                        INFO "swapD: pickup devices from systemd-gpt-auto-generator"
                        FIND_SWAP_UNITS | while read -r UNIT_PATH; do
                                if grep -q systemd-gpt-auto-generator $UNIT_PATH; then
                                        DEV=$(GET_WHAT_FROM_SWAP_UNIT $UNIT_PATH)
                                        UNIT_NAME=$(basename $UNIT_PATH)
                                        swapoff "$DEV"
                                        rm -f "$UNIT_NAME"
                                fi
                        done

                        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
                                UNIT_NAME=$(gen_swap_unit What="$device" Options=discard Tag="swapd")
                                systemctl daemon-reload
                                systemctl start $UNIT_NAME || continue
                                INFO "swapD: enabled device: $device"
                                swapd_prio=$((swapd_prio-1))
                        done
                fi
        ;;
        stop)
                [ ! -f $LOCK_STARTED ] && INFO_E "$0 already stopped"

                FIND_SWAP_UNITS | while read -r UNIT_PATH; do
                        if grep -q swapd $UNIT_PATH; then
                                DEV=$(GET_WHAT_FROM_SWAP_UNIT $UNIT_PATH)
                                UNIT_NAME=$(basename $UNIT_PATH)
                                swapoff "$DEV"
                                rm -f "$UNIT_NAME"
                        fi
                done

                if [ -f "$WORK_DIR/swapfc/.lock" ]; then
                        rm -f "$WORK_DIR/swapfc/.lock"
                        sleep 5 # Wait some time
                fi

                FIND_SWAP_UNITS | while read -r UNIT_PATH; do
                        if grep -q swapfc $UNIT_PATH; then
                                DEV=$(GET_WHAT_FROM_SWAP_UNIT $UNIT_PATH)
                                UNIT_NAME=$(basename $UNIT_PATH)
                                swapoff "$DEV"
                                rm -f "$UNIT_NAME"
                                [ -f $DEV ] && rm -f $DEV
                        fi
                done

                FIND_SWAP_UNITS | while read -r UNIT_PATH; do
                        if grep -q zram $UNIT_PATH; then
                                DEV=$(GET_WHAT_FROM_SWAP_UNIT $UNIT_PATH)
                                UNIT_NAME=$(basename $UNIT_PATH)
                                INFO "Zram: swapoff $DEV"
                                swapoff "$DEV"
                                rm -f $UNIT_PATH
                                zramctl -r "$DEV"
                        fi
                done

                if [ -d $WORK_DIR/zswap/ ]; then
                        INFO "Zswap: restore configuration: start"
                        for file in $WORK_DIR/zswap/*; do
                                read_l "$file" > "$ZSWAP_M_P/$(basename $file)"
                        done
                        INFO "Zswap: restore configuration: complete"
                fi

                wait

                rm -f $LOCK_STARTED
        ;;
        status)
                SWAP_TOTAL=$(grep SwapTotal: /proc/meminfo | grep -o '[0-9]*')
                SWAP_FREE=$(grep SwapFree: /proc/meminfo | grep -o '[0-9]*')
                SWAP_USED=$(((SWAP_TOTAL-SWAP_FREE)*1024))

                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))

                        RATIO=0
                        ((STORED_PAGES>0)) && RATIO=$((USED_PAGES*100/STORED_PAGES))

                        {
                                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}%
                                ((SWAP_USED > 0)) && \
                                        echo . . zswap_store/swap_store ${STORED_BYTES}/${SWAP_USED} $((STORED_BYTES*100/SWAP_USED))%
                        } | 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/swapfc ]; then
                        {
                                echo SwapFC:
                                swapon --show TYPE | grep -e 'NAME\|file\|loop' | sed 's/^/. /g'
                        } | column -t
                fi
        ;;
        *) help ;;
esac
