#!/bin/bash

# ----------------------------------------------------------------------------
# zramen: manage zram swap space
# ----------------------------------------------------------------------------

# ==============================================================================
# constants {{{

# use this compression algorithm for zram by default
readonly COMP_ALGORITHM='lz4'
readonly ZRAM_COMP_ALGORITHM="${ZRAM_COMP_ALGORITHM:-$COMP_ALGORITHM}"

# give zram swap device highest priority
readonly PRIORITY=32767
readonly ZRAM_PRIORITY="${ZRAM_PRIORITY:-$PRIORITY}"

# allocate this percentage of memory for zram by default
readonly SIZE=25
readonly ZRAM_SIZE="${ZRAM_SIZE:-$SIZE}"

# set the maximum number of compression streams for zram
readonly STREAMS=$(nproc)
readonly ZRAM_STREAMS="${ZRAM_STREAMS:-$STREAMS}"

# zramen version number
readonly VERSION=0.2.1

# set TMPDIR to work directory for mktemp
readonly WORK_DIR='/var/run/zramen'
TMPDIR="$WORK_DIR"

# end constants }}}
# ==============================================================================
# usage {{{

_usage() {
read -r -d '' _usage_string <<EOF
Usage:
  zramen [-h|--help] <command>
  zramen [-a|--algorithm <algo>]
         [-n|--num <uint>]
         [-s|--size <uint>]
         [-p|--priority <int>]
         make
  zramen toss

Options:
  -h, --help       Show this help text
  -v, --version    Show program version
  -a, --algorithm  Compression algorithm for zram (Default: $ZRAM_COMP_ALGORITHM)
  -n, --num        Number of compression streams for zram (Default: $ZRAM_STREAMS)
  -p, --priority   Priority of zram swap device (Default: $ZRAM_PRIORITY)
  -s, --size       Percentage of memory to allocate for zram (Default: $ZRAM_SIZE)

Commands:
  make        Make zram swap device
  toss        Remove zram swap device

Algorithm
  Run zramctl --help to see a list of acceptable algorithms:
  | lzo
  | lzo-rle
  | lz4
  | lz4hc
  | zstd
  | deflate
  | 842

Num
  Number of zram compression streams; try one per core

Priority
  Must be an integer <= 32767; higher number means higher zram priority

Size
  Percentage of memory to allocate for zram; try <= 50
EOF
echo "$_usage_string"
}

_POSITIONAL=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    -h|--help)
      _usage
      exit 0
      ;;
    -v|--version)
      echo "$VERSION"
      exit 0
      ;;
    -a|--algorithm)
      _algorithm="$2"
      shift
      shift
      ;;
    -n|--num)
      _streams="$2"
      # shift past argument and value
      shift
      shift
      ;;
    -p|--priority)
      _priority="$2"
      shift
      shift
      ;;
    -s|--size)
      _size="$2"
      shift
      shift
      ;;
    -*)
      # unknown option
      _usage
      exit 1
      ;;
    make|toss)
      _POSITIONAL+=("$1")
      shift
      ;;
    *)
      # unknown command
      _usage
      exit 1
      ;;
  esac
done

if ! [[ "${#_POSITIONAL[@]}" == '1' ]]; then
  _usage
  exit 1
fi

# restore positional params
set -- "${_POSITIONAL[@]}"

# end usage }}}
# ==============================================================================

# sanitize compression algorithm input
_algorithm="${_algorithm:-$ZRAM_COMP_ALGORITHM}"
case "$_algorithm" in
  # proper algo chosen, no action necessary
  lzo|lzo-rle|lz4|lz4hc|zstd|deflate|842)
    ;;
  # improper algo chosen, reset to default
  *)
    _algorithm="$COMP_ALGORITHM"
    ;;
esac

# sanitize priority input
_priority=${_priority:-$ZRAM_PRIORITY}
[[ $_priority -le 32767 ]] \
  || _priority=32767

# sanitize size input
_size=${_size:-$ZRAM_SIZE}
[[ $_size -gt 0 ]] \
  || _size=$SIZE
[[ $_size -le 100 ]] \
  || _size=100

# sanitize streams input
_streams=${_streams:-$ZRAM_STREAMS}
# number of CPUs requested must be 1+
[[ $_streams -gt 0 ]] \
  || _streams=1
# number of CPUs requested must not exceed CPUs available
[[ $_streams -le $STREAMS ]] \
  || _streams=$STREAMS

INFO() {
  echo "zramen#info: $*"
}

WARN() {
  echo "zramen#warn: $*"
}

ERRO() {
  echo "zramen#erro: $*"
  exit 1
}

calc() {
  # truncate to whole number
  printf '%.f' "$(bc --mathlib <<< "$@")"
}

make() {
  local _mem_total
  local _mem_to_alloc
  local _zram_dev

  _mem_total=$(grep 'MemTotal:' /proc/meminfo | awk '{print $2}')
  _mem_to_alloc=$(calc "$_mem_total * 1024 * ($_size / 100)")

  if ! [[ -d '/sys/module/zram' ]]; then
    INFO 'Attempting to find zram module - not part of kernel'
    modprobe --dry-run zram 2>/dev/null \
      || ERRO 'Sorry, could not find zram module'
    # loop to handle 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 'Attempting to initialize free device'
    _tmp="$(mktemp)"
    # return name of first free device
    zramctl \
      --algorithm "$_algorithm" \
      --find \
      --size "$_mem_to_alloc" \
      --streams "$_streams" &> \
      "$_tmp"
    read -r _output < "$_tmp"
    rm -f "$_tmp"
    unset _tmp
    case "$_output" in
      *'failed to reset: Device or resource busy'*)
        sleep 1
        ;;
      *'zramctl: no free zram device found'*)
        WARN 'zramctl could not find free device'
        INFO 'Attempting zram hot add'
        ! [[ -f '/sys/class/zram-control/hot_add' ]] \
          && ERRO 'Sorry, this kernel does not support zram hot add'
        read -r _hot_add < /sys/class/zram-control/hot_add
        INFO "Hot added new zram swap device: /dev/zram$_hot_add"
        ;;
      /dev/zram*)
        [[ -b "$_output" ]] \
          || continue
        _zram_dev="$_output"
        break
        ;;
    esac
  done

  if [[ -b "$_zram_dev" ]]; then
    INFO "Successfully initialized zram swap device: $_zram_dev"
    mkdir -p "$WORK_DIR/zram"
    mkswap "$_zram_dev" --label "$(basename "$_zram_dev")" &> /dev/null \
      && swapon --discard --priority $_priority "$_zram_dev" \
      && ln --symbolic "$_zram_dev" "$WORK_DIR/zram/"
  else
    WARN 'Could not get free zram device'
  fi
}

toss() {
  for zram in "$WORK_DIR/zram"/*; do
    ! [[ -b $zram ]] \
      && continue
    INFO "Removing zram swap device: /dev/$(basename "$zram")"
    swapoff "$zram" \
      && zramctl --reset "$(basename "$zram")" \
      && rm "$zram" \
      && INFO "Removed zram swap device: /dev/$(basename "$zram")"
  done
  [[ -d "$WORK_DIR" ]] \
    && rm -rf "$WORK_DIR"
}

main() {
  if ! [[ "$UID" == '0' ]]; then
    echo 'Sorry, requires root privileges'
    exit 1
  fi
  [[ -d "$WORK_DIR" ]] \
    || mkdir -p "$WORK_DIR"
  if [[ "$1" == 'make' ]]; then
    make
  elif [[ "$1" == 'toss' ]]; then
    toss
  else
    # unknown command
    _usage
    exit 1
  fi
}

main "$1"

# vim: set filetype=sh foldmethod=marker foldlevel=0 nowrap:
