#!/usr/bin/env sh
# (C) Martin V\"ath <martin@mvath.de>"
set -u
Echo() {
	printf '%s\n' "$*"
}

Echon() {
	printf '%s' "$*"
}

Warning() {
	Echo "${0##*/}: $*" >&2
}

Fatal() {
	Warning "$*"
	exit 1
}

Usage() {
	Echo "Usage: ${0##*/} [options] SIZE [DIR]
Prepare a zram device and use it as swap (resp. mount it under DIR).
SIZE is the maximal size in megabytes.
For umounting/freeing the zram device, use SIZE=0.
If DIR is - then only a filesystem is created in /dev/zram\$DEV (or the device
is removed if SIZE is 0) but it is not mounted
(options -o -c -m -T take no effect in this case, of course).
The latter can be useful e.g. for btrfs if multiple devices should be mounted
jointly afterwards.
The following options are available.
An empty argument means the same as if the option is not specified.
-d DEV   Use zram device DEV. If not specified, DEV=0 is assumed.
         Make sure to use the matching value for umounting (SIZE=0)!
-D NUM   If modprobe needs to be used, require NUM devices. This is not
         recommended. Rely instead on /etc/modprobe.d/zram.conf with the line
         options zram num_devices=NUM
-s NUM   Use up to NUM parallel compression streams for the device
-a ALGO  Set compression algorithm to ALGO (e.g. lz4 or lzo)
-c OWNER If specified, chown to OWNER (resp. OWNER:GROUP) after mounting.
         If not specified, the default owner (root:root) is not changed.
-m MODE  If specified, chmod DIR to MODE after mounting.
-o OPTS  If specified, mount DIR with -o OPTS.
-p PRIO  Use priority PRIO for the swap device.
         If not specified, PRIO=16383 is assumed.
         Use PRIO=- if you want to keep the default priority (-1).
-t TYPE  Use a filesystem of type TYPE if DIR is specified.
         TYPE can be either ext2, ext4, or btrfs
         If not specified, TYPE=ext4 is assumed.
-i IRATIO If specified override default bytes/inodes in fs (ext2, ext4)
-N INODES If specified override inode count (ext2, ext4)
-T       If specified, do not use the discard (TRIM) feature of ext4.
         Use this option with linux-3.14 or earlier or when you want a slight
         speed increase at the cost of possibly wasting a lot of memory
-L       Do not use zramctl even if available
-k       Do no attempt to umount/free a previously used zram under this device."
	exit ${1:-1}
}

PATH=${PATH-}${PATH:+:}/usr/sbin:/sbin
dev=
num=
streams=
algo=
prio=
fstype=
opts=
mode=
owgr=
iratio=
inodes=
keep=false
zramctl=:
discard=:
OPTIND=1
while getopts 's:a:c:d:D:m:o:p:t:i:N:TLkhH' opt
do	case $opt in
	s)	streams=$OPTARG;;
	a)	algo=$OPTARG;;
	c)	owgr=$OPTARG;;
	d)	dev=$OPTARG;;
	D)	num=$OPTARG;;
	m)	mode=$OPTARG;;
	o)	opts=$OPTARG;;
	p)	prio=$OPTARG;;
	t)	fstype=$OPTARG;;
	i)	iratio=$OPTARG;;
	N)	inodes=$OPTARG;;
	T)	discard=false;;
	L)	zramctl=false;;
	k)	keep=:;;
	'?')	exit 1;;
	*)	Usage 0;;
	esac
done
shift $(( $OPTIND -1 ))
[ $# -eq 1 ] || [ $# -eq 2 ] || Usage

if $zramctl
then	command -v zramctl >/dev/null 2>&1 || zramctl=false
fi

# Defaults for empty/unspecified options:
: ${dev:=0} ${prio:=16383} ${fstype:=ext4}

umount=:
[ "${1:-0}" -ge 0 ] || Usage
if [ "${1:-0}" -ne 0 ]
then	umount=false
	size=$(( $1 * 1024 * 1024 )) && [ "$size" -gt 0 ] || Usage
fi

swap=:
if [ $# -eq 2 ]
then	swap=false
	mount=:
	dir=${2%/}
	if [ x"$dir" = x'-' ]
	then	mount=false
	elif [ -z "$dir" ] || ! test -d "$dir"
	then	if $umount
		then	Warning "${dir:-(empty)} is not a directory. Umounting anyway"
		else	Fatal "${dir:-(empty)} is not a directory"
		fi
	fi
fi

devnode=/dev/zram$dev
if ! test -e "$devnode"
then	$umount && exit
	modprobe zram ${num:+"num_devices=$num"} >/dev/null 2>&1
	test -e "$devnode" || Fatal "cannot find $devnode"
fi
if ! $zramctl
then	block=/sys/block/zram$dev
	test -d "$block" || Fatal "cannot find $block"
fi

if ! $keep
then	if $swap
	then	swapoff -- "$devnode" >/dev/null 2>&1
	else	! $mount || umount -- "$devnode" >/dev/null 2>&1
	fi && {
		if $zramctl
		then	zramctl --reset -- "$devnode"
		else	Echon 1 >"$block/reset"
		fi || Warning "warning: failed to reset zram$dev"
	}
fi
status=$?
$umount && exit $status

if $zramctl
then	zramctl --size "$size" \
		${streams:+--streams "$streams"} \
		${algo:+--algorithm "$algo"} \
		-- "$devnode" || Fatal "zramctl zram$dev failed"
else	[ -z "$streams" ] || Echon "$streams" >"$block/max_comp_streams" || \
		Warning "warning: failed to set max_comp_streams to $streams"
	[ -z "$algo" ] || Echon "$algo" >"$block/comp_algorithm" || \
		Warning "warning: failed to set comp_algorithm to $algo"
	Echon "$size" >"$block/disksize" || Fatal "cannot set zram$dev size"
fi

if $swap
then	mkswap "$devnode" >/dev/null || Fatal "mkswap $devnode failed"
	if [ x"$prio" = x'-' ]
	then	swapon -- "$devnode"
	else	swapon -p "$prio" -- "$devnode"
	fi || Fatal "swapon failed"
	exit 0
fi

fsopts='-O^huge_file,sparse_super'
case $fstype in
ext2)
	mkfs.ext2 -m0 "$fsopts" ${iratio:+-i "$iratio"} \
		${inodes:+-N "$inodes"} "$devnode" >/dev/null \
		|| Fatal "mkfs.ext2 $devnode failed"
	tune2fs -c0 -i0 -m0 "$devnode" >/dev/null
	$mount || exit 0
	mount -t ext2 ${opts:+-o "$opts"} -- "$devnode" "$dir" \
		|| Fatal "mount $devnode failed";;
ext4)
	fsopts=$fsopts',extent,^uninit_bg,dir_nlink,extra_isize,^has_journal'
	mkfs.ext4 -m0 "$fsopts" ${iratio:+-i "$iratio"} \
		${inodes:+-N "$inodes"} "$devnode" >/dev/null \
		|| Fatal "mkfs.ext4 $devnode failed"
	tune2fs -c0 -i0 -m0 "$devnode" >/dev/null
	$mount || exit 0
	! $discard || opts=$opts${opts:+,}'discard'
	mount -t ext4 ${opts:+-o "$opts"} -- "$devnode" "$dir" \
		|| Fatal "mount $devnode failed";;
btrfs)
	mkfs.btrfs -d single -m single "$devnode" >/dev/null \
		|| Fatal "mkfs.btrfs $extnode failed"
	$mount || exit 0
	mount -t btrfs ${opts:+-o "$opts"} -- "$devnode" "$dir" \
		|| Fatal "mount $devnode failed";;
*)
	Fatal "unsupported filesystem $fstype";;
esac

# Change owner/mode if requested
[ -z "$owgr" ] || chown -- "$owgr" "$dir"
[ -z "$mode" ] || chmod -- "$mode" "$dir"
