zram-init

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

SysCtl() {
	printf '%s' "$1" >|"$2"
}

SysCtlN() {
	printf '%s\n' "$1" >|"$2"
}

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

Warning() {
	Error "warning: $*" >&2
}

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

Push() {
	PushA_=`push.sh 2>/dev/null` || Fatal \
'push.sh from https://github.com/vaeth/push (v2.0 or newer) required'
	eval "$PushA_"
	Push "$@"
}

IsNumeric() {
	case ${1:-x} in
	*[!01234567890]*)
		return 1;;
	esac
	:
}

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
-S MAX   Use maximally MAX megabytes of uncompressed memory for the device
-b DEV   Use DEV as backing 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)
-L LABEL Specify the label LABEL for the new filesystem
-U UUID  Specify the uuid UUID for the new filesystem
-T       If specified, do not use the discard (TRIM) feature of ext4/swap.
         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

If you have push.sh in \$PATH, you can also use accumulatively:
-K ARG   Pass ARG to the respective mkswap or mkfs.* call
-M ARG   Pass ARG to the respective swapon/mount call
-2 ARG   Pass ARG to the tune2fs call (ignored unless for ext2 or ext4)
-Z ARG   Pass ARG to the zramctl call"
	exit ${1:-1}
}

PATH=${PATH-}${PATH:+:}/usr/sbin:/sbin
dev=
num=
streams=
algo=
prio=
fstype=
opts=
mode=
owgr=
iratio=
inodes=
label=
uuid=
keep=false
zramctl=:
discard=:
mkfs_opt=
mount_opt=
tune2fs_opt=
zramctl_opt=
mem_limit=
backing_dev=
OPTIND=1
while getopts 's:S:a:b:c:d:D:m:o:p:t:i:N:L:U:TlkK:M:2:Z:hH' opt
do	case $opt in
	s)	streams=$OPTARG;;
	S)	mem_limit=$OPTARG;;
	a)	algo=$OPTARG;;
	b)	backing_dev=$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;;
	L)	label=$OPTARG;;
	U)	uuid=$OPTARG;;
	T)	discard=false;;
	l)	zramctl=false;;
	k)	keep=:;;
	K)	Push mkfs_opt "$OPTARG";;
	M)	Push mount_opt "$OPTARG";;
	2)	Push tune2fs_opt "$OPTARG";;
	Z)	Push zramctl_opt "$OPTARG";;
	'?')	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}

IsNumeric "$dev" || Fatal "device $dev is not numeric"
IsNumeric "${1:-0}" || Fatal "size $1 is not numeric"

umount=:
if [ ${1:-0} -ne 0 ]
then	umount=false
	size=$(( $1 * 1024 * 1024 )) && [ $size -gt 0 ] || \
		Fatal 'failed to calculate size'
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
zclass=/sys/class/zram-control

HotAdd() {
	while :
	do	curradd=
		exec 3<&0 && {
			exec <"$zclass/hot_add" && \
				read curradd || curradd=
			exec 0<&3
			exec 3<&-
		}
		IsNumeric "$curradd" || {
			Warning "hot_add failed for $devnode"
			break
		}
		[ $curradd -lt $dev ] || break
	done
	sleep=0
	while ! test -b "$devnode"
	do	sleep 1
		sleep=$(( $sleep + 1 ))
		[ $sleep -lt 5 ] || Fatal "failed to create $devnode"
	done
}

if ! test -b "$devnode"
then	$umount && exit
	IsNumeric "${num:-0}" || Fatal "device count $num is not numeric"
	modprobe zram ${num:+num_devices=$num} >/dev/null 2>&1
	sleep=0
	while ! test -b "$devnode" && ! test -d "$zclass"
	do	sleep 1
		sleep=$(( $sleep + 1 ))
		[ $sleep -lt 5 ] || \
			Fatal "cannot create $devnode: $zclass missing"
	done
	test -b "$devnode" || {
		HotAdd
		keep=:
	}
fi
block=/sys/block/zram$dev
if ! $zramctl || [ -n "$mem_limit" ] || [ -n "${backing_dev:++}" ]
then	test -d "$block" || Fatal "cannot find $block"
fi

status=0
if ! $keep
then	if $swap
	then	swapoff -- "$devnode" >/dev/null 2>&1 || status=$?
	elif $mount
	then	umount -- "$devnode" >/dev/null 2>&1 || status=$?
	fi
	if [ $status -eq 0 ]
	then	if $zramctl
		then	zramctl --reset -- "$devnode" || status=$?
		else	SysCtl 1 "$block/reset" || status=$?
			if $umount && test -w "$zclass/hot_remove"
			then	SysCtl "$device" "$zclass/hot_remove" \
					|| status=$?
			fi
		fi
		[ $status -eq 0 ] || \
			Warning "failed to reset zram$dev"
	fi
fi
$umount && exit $status
test -b "$devnode" || HotAdd

if $zramctl && [ -z "${backing_dev:++}" ]
then	eval "set -- a $zramctl_opt"
	shift
	zramctl --size "$size" \
		${streams:+--streams "$streams"} \
		${algo:+--algorithm "$algo"} ${1+"$@"} \
		-- "$devnode" || Fatal "zramctl zram$dev failed"
else	[ -z "$streams" ] || SysCtl "$streams" "$block/max_comp_streams" || \
		Warning "failed to set zram$dev max_comp_streams to $streams"
	[ -z "$algo" ] || SysCtl "$algo" "$block/comp_algorithm" || \
		Warning "failed to set zram$dev comp_algorithm to $algo"
	[ -z "${backing_dev:++}" ] || \
		SysCtlN "$backing_dev" "$block/backing_dev" || \
		Warning "failed to set zram$dev backing_dev"
	SysCtl "$size" "$block/disksize" || Fatal "cannot set zram$dev size"
fi
[ -z "$mem_limit" ] || SysCtl "$mem_limit" "$block/mem_limit" || \
	Warning "failed to set zram$dev mem_limit"

eval "set -- a $mkfs_opt"
shift

if $swap
then	mkswap ${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null || Fatal "mkswap $devnode failed"
	discardopt=-d
	$discard || discardopt=
	[ x"$prio" != x'-' ] || prio=
	swapon ${prio:+-p "$prio"} $discardopt -- "$devnode" \
		|| Fatal "swapon $devnode failed"
	exit 0
fi

fsopts='-O^huge_file,sparse_super'
case $fstype in
ext2)
	mkfs.ext2 -m0 "$fsopts" ${iratio:+-i "$iratio"} \
		${inodes:+-N "$inodes"} \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| Fatal "mkfs.ext2 $devnode failed"
	eval "set -- a $tune2fs_opt"
	shift
	tune2fs -c0 -i0 -m0 -- "$devnode" >/dev/null
	$mount || exit 0
	eval "set -- a $mount_opt"
	shift
	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"} \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| Fatal "mkfs.ext4 $devnode failed"
	eval "set -- a $tune2fs_opt"
	shift
	tune2fs -c0 -i0 -m0 -- "$devnode" >/dev/null
	$mount || exit 0
	! $discard || opts=$opts${opts:+,}'discard'
	eval "set -- a $mount_opt"
	shift
	mount -t ext4 ${opts:+-o "$opts"} -- "$devnode" "$dir" \
		|| Fatal "mount $devnode failed";;
btrfs)
	mkfs.btrfs -d single -m single \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| Fatal "mkfs.btrfs $extnode failed"
	$mount || exit 0
	eval "set -- a $mount_opt"
	shift
	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"