#!/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}"

# set zramen log level
readonly QUIET=0
readonly ZRAMEN_QUIET="${ZRAMEN_QUIET:-$QUIET}"

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

# set maximum size of zram in MiB
readonly MAX_SIZE=4096

# zramen version number
readonly VERSION=1.0.0

# 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>]
         [-m|--max-size <uint>]
         [-p|--priority <int>]
         [-b|--backing-device <dev>]
         [-q|--quiet]
         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)
  -p, --priority         Priority of zram swap device (Default: $ZRAM_PRIORITY)
  -s, --size             Percentage of memory to allocate for zram (Default: $ZRAM_SIZE)
  -m, --max-size         Maximum size of zram in MiB (optional)
  -b, --backing-device   Backing device to use for incompressible pages
  -q, --quiet            Disable info messages

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

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

Size
  Percentage of memory to allocate for zram; try <= 50

Max Size
  Maximum size of created zram in MiB; use if plenty of memory

Backing Device
  Backing device to use for incompressible pages

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
      ;;
    -p|--priority)
      _priority="$2"
      shift
      shift
      ;;
    -s|--size)
      _size="$2"
      shift
      shift
      ;;
    -m|--max-size)
      _max_size="$2"
      shift
      shift
      ;;
    -b|--backing-device)
      _backing_device="$2"
      shift
      shift
      ;;
    -q|--quiet)
      _quiet=1
      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 }}}
# ==============================================================================

_quiet="${_quiet:-$ZRAMEN_QUIET}"
case "$_quiet" in
  0|1)
    ;;
  *)
    WARN 'Improper value for quiet, using default'
    _quiet=0
    ;;
esac

# 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)
    ;;
  *)
    WARN 'Improper compression algorithm chosen, using 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 250 ]] \
  || _size=250

# sanitize max size input and mark seen if applicable
if [[ -n $_max_size ]] || [[ -n $ZRAM_MAX_SIZE ]]; then
  _max_size_set=true
  _max_size=${_max_size:-$ZRAM_MAX_SIZE}
  [[ $_max_size -gt 0 ]] \
    || _max_size=$MAX_SIZE
fi

INFO() {
  if [[ $_quiet -eq 0 ]] then
    echo "zramen#info: $*"
  fi
}

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

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

make() {
  local _mem_total
  local _mem_to_alloc
  local _zram_dev

  _mem_total=$(awk '/MemTotal:/ {printf("%.0f",$2/1024)}' /proc/meminfo)
  _mem_to_alloc=$(awk -v size="$_size" '{printf("%.0f",$0*(size/100))}' <<< "$_mem_total")

  # only enforce max size when specified
  if [[ -n $_max_size_set ]]; then
    if [[ $_mem_to_alloc -gt $_max_size ]]; then
      _mem_to_alloc=$_max_size
    fi
  fi

  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'
    _output="$(zramctl --find 2>&1)"
    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"

    if [[ -n "$_backing_device" ]]; then
      if [[ -b "$_backing_device" ]]; then
        echo 1 > /sys/block/$(basename "$_zram_dev")/reset \
          && echo "$_backing_device" > /sys/block/$(basename "$_zram_dev")/backing_dev \
          && INFO "Successfully initialized zram backing device: $_backing_device" \
          || WARN "Failed to write to /sys/block/$(basename "$_zram_dev")/backing_dev"
      else
        WARN "Failed to find backing device $_backing_device"
      fi
    else
      INFO "Continuing without any backing device"
    fi

    zramctl \
      "$_zram_dev" \
      --algorithm "$_algorithm" \
      --size "${_mem_to_alloc}MiB"

    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:
