#!/bin/bash

set -euo pipefail

size=1M
threshold=1000
loss=1000
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]
	-l <loss>		PSI integral drop per second.
				[default: $loss]
	-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"
	integral=0

	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 and loss by interval
	threshold="$((threshold * interval))"
	loss="$((loss * interval))"

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

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

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

	if [[ $integral -lt $loss ]]; then
		integral=0
	else
		integral=$((integral - loss))
	fi

	integral=$((integral + delta))
}

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

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

	return 0
}

calc_size_to_reclaim_kb() {
	# Rescale reclaim target by integral
	_diff=$((threshold - integral))
	_size_kb=$((size_kb * _diff / threshold))

	echo $_size_kb
}

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

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

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

	log_debug "memusage: $_usage"
}

wait_for_next_run() {
	sleep "$interval"
}

start_reclaim() {
	reclaim_size_kb="$(calc_size_to_reclaim_kb)"

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

	log_info "reclaimed memory ${reclaim_size_kb}K"
}

reclaim_run() {
	update_states

	[ -n "$debug" ] && dump_reclaim_args

	if should_reclaim; then
		start_reclaim
	fi

	[ -n "$debug" ] && dump_memusage
}

while getopts "s:t:i:l:dh" arg; do
	case "$arg" in
		s) size="$OPTARG" ;;
		t) threshold="$OPTARG" ;;
		i) interval="$OPTARG" ;;
		l) loss="$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