#!/sbin/openrc-run
# SPDX-License-Identifier: MIT

description="Create swap on zram device"

depend() {
	after bootmisc modules
}

fail() {
	eend 1 "$1"
	exit 1
}

numfmt_to_bytes() {
	# remove B suffix if any
	local size="${1%[bB]}" unit

	case "$size" in
	*[!0-9]*[0-9]|*[!0-9]*[a-zA-Z]|[!0-9])
		fail "Unexpected size format $size (expected e.g. 300M)";;
	*[0-9]) echo "$size";;
	esac

	unit="${size##*[0-9]}"
	size="${size%[a-zA-Z]}"
	case "$unit" in
	[Kk]) echo "$((size * 1000))";;
	[Mm]) echo "$((size * 1000000))";;
	[Gg]) echo "$((size * 1000000000))";;
	*) fail "unsupported unit in $size (only support G/M/K)";;
	esac
}

get_size() {
	local memtotal
	# prefer configured size if any
	[ -n "$size" ] && return

	memtotal=$(awk '/MemTotal/ { print $2 }' /proc/meminfo)
	if [ "$memtotal" -lt 393216 ]; then
		size=0
	elif [ "$memtotal" -lt 774144 ]; then
		size=256M
	elif [ "$memtotal" -lt 1572864 ]; then
		size=512M
	else
		size=1G
	fi
}

zram_find() {
	local dev
	for dev in /sys/block/zram*; do
		[ -e "$dev" ] || continue
		if [ "$(cat "$dev/disksize")" = 0 ]; then
			echo "${dev#/sys/block/}"
			return
		fi
	done
	[ -e /sys/class/zram-control/hot_add ] \
		|| fail "zram support not compiled in?"
	dev="zram$(cat /sys/class/zram-control/hot_add)" \
		|| fail "Could not hot-add new zram device"
	echo "$dev"
}

zram_remove() {
	local dev="$1"
	echo "${dev##*zram}" > /sys/class/zram-control/hot_remove
}

zram_set() {
	local dev="/sys/block/$1" size="$2" algo="$3"

	if [ -n "$extraOpts" ]; then
		ewarn "Ignoring extra zramctl options $extraOpts"
	fi
	echo 1 > "$dev/reset" || fail "Could not reset $1"
	if [ -n "$algo" ]; then
		echo "$algo" > "$dev/comp_algorithm" \
			|| fail "Could not set $1 compression algorithm to $algo"
	fi
	numfmt_to_bytes "$size" > "$dev/disksize" \
		|| fail "Could not set $1 size to $size"
}

start() {
	local dev

	get_size
	[ "$size" = 0 ] && return

	ebegin "Creating zram swap device"

	dev=$(zram_find)
	zram_set "$dev" "$size" "${algo:-zstd}"

	if ! mkswap "/dev/$dev" >/dev/null; then
		eend $? "Could not format zram"
		zram_remove "$dev"
		return 1
	fi

	if ! swapon -p "${priority:-100}" "/dev/$dev"; then
		eend $? "Could not activate swap"
		zram_remove "$dev"
		return 1
	fi

	eend 0
}

stop() {
	local dev
	local rc=0

	ebegin "Deactivating zram swap device"
	while read -r dev _; do
		[ "${dev#/dev/zram}" != "$dev" ] || continue

		if ! swapoff "$dev"; then
			eerror "Could not deactivate swap"
			rc=1
			continue
		fi
		if ! zram_remove "$dev"; then
			eerror "Could not reset zram device"
			rc=1
		fi
	done < /proc/swaps

	eend "$rc" "Could not reset all zram swap devices"
}
