#!/bin/bash

set -euo pipefail

size=1M
threshold=100
interval=1
debug=

usage() {
	cat <<EOF
Usage: $0 [OPTIONS]

Reclaim memory at fixed interval

Options:
	-s <size>		Memory to reclaim in one run. Available units: K, M
				[default: $size]
	-t <threshold>		PSI threshold to continue reclaiming.
				[default: $threshold]
	-i <interval>		Interval between reclaims in seconds.
				[default: $interval]
	-d			Enable debug output.
				[default: no]
	-h			Show this message.
EOF
	exit "$1"
}

log() {
	printf "[%s] %5s: %s\n" "$(date +%F\ %T)" "$1" "$2" >&2
}

log_debug() {
	log "DEBUG" "$1"
}

log_info() {
	log "INFO" "$1"
}

log_warn() {
	log "WARN" "$1"
}

log_fatal() {
	log "FATAL" "$1"
}

die() {
	log_fatal "$1"
	exit 1
}

readpsi() {
	psi_path="/proc/pressure/memory"
	head -n 1 "$psi_path" | awk -F= '{ print $5 }'
}

init_reclaim_states() {
	last_psi="$(readpsi)"
	psi="$last_psi"

	size_kb=
	# Convert size into KBs
	case "$size" in
		[0-9]*M|[0-9]*m)
			size_kb="$(echo "$size" | tr -d mM)"
			size_kb="$((size_kb * 1024))"
			;;
		[0-9]*K|[0-9]*k)
			size_kb="$(echo "$size" | tr -d kK)"
			;;
		*)
			die "invalid size format: $size"
			;;
	esac

	# Scale PSI threshold by interval
	threshold="$((threshold * interval))"

	log_info "reclaim params:"
	log_info "size=$size threshold=$threshold interval=$interval"
}

update_states() {
	last_psi="$psi"
	psi="$(readpsi)"

	delta="$((psi - last_psi))"
}

skip_reclaim() {
	log_info "reclaim skipped for $1"
}

should_reclaim() {
	if [[ "$delta" -gt "$threshold" ]]; then
		skip_reclaim "high pressure: $delta > $threshold over the last ${interval}s"
		return 1
	fi

	return 0
}

reclaim() {
	cgpath="/sys/fs/cgroup"
	echo "${size_kb}K" > "$cgpath/memory.reclaim"
}

dump_reclaim_args() {
	log_debug "last_psi=$last_psi psi=$psi delta=$delta"
}

dump_memusage() {
	NL=$'\n'
	_usage="$NL$(free -h)$NL"

	log_debug "memusage: $_usage"
}

wait_for_next_run() {
	sleep "$interval"
}

reclaim_run() {
	update_states

	[ -n "$debug" ] && dump_reclaim_args

	if ! should_reclaim; then
		return
	fi

	if ! reclaim; then
		log_warn "unable to reclaim, status=$?"
		return
	fi

	log_info "reclaimed memory ${size_kb}K"

	[ -n "$debug" ] && dump_memusage

	return
}

while getopts "s:t:i:dh" arg; do
	case "$arg" in
		s) size="$OPTARG" ;;
		t) threshold="$OPTARG" ;;
		i) interval="$OPTARG" ;;
		d) debug=y ;;
		h) usage 0 ;;
		?) usage 1 ;;
	esac
done

shift $((OPTIND - 1))

init_reclaim_states

while true; do
	reclaim_run
	wait_for_next_run
done