#!/usr/bin/env bash
# Run a heavy command on the 8 GB Pi so that running out of memory kills only this command, never the desktop session.
#   MEM=4G MIN_FREE_MB=1200 bash scripts/pi_run.sh <command...>
# 1. When the memory cgroup is available (cmdline.txt: cgroup_enable=memory cgroup_memory=1), the command runs in a
#    systemd user scope with MemoryMax=$MEM and no swap: the kernel kills it alone when it exceeds the cap.
# 2. In every case a guard stops the command's process group when the system's MemAvailable drops below
#    MIN_FREE_MB (covers other processes filling memory, and systems without the memory cgroup).
# Prints the peak memory of the command when it ends.
set -uo pipefail
MEM=${MEM:-4G}
MIN_FREE_MB=${MIN_FREE_MB:-1200}
uid=$(id -u)
unit="simurg-job-$$-$(date +%s)"
cg="/sys/fs/cgroup/user.slice/user-$uid.slice/user@$uid.service/app.slice/$unit.scope"
if grep -qw memory "/sys/fs/cgroup/user.slice/user-$uid.slice/user@$uid.service/cgroup.controllers" 2>/dev/null; then
  setsid systemd-run --user --scope --unit="$unit" -p MemoryMax="$MEM" -p MemorySwapMax=0 --quiet -- "$@" &
  mode="cgroup MemoryMax=$MEM"
else
  setsid "$@" &
  mode="no memory cgroup"
fi
pid=$!
peak=0
tripped=0
while kill -0 "$pid" 2>/dev/null; do
  if [[ -r $cg/memory.peak ]]; then
    peak=$(( $(cat "$cg/memory.peak") / 1024 ))
  else
    rss=$(ps -o rss= -g "$pid" 2>/dev/null | awk '{s+=$1} END {print s+0}'); (( rss > peak )) && peak=$rss
  fi
  avail=$(awk '/^MemAvailable:/ {print int($2/1024)}' /proc/meminfo)
  if (( avail < MIN_FREE_MB )); then
    echo "pi_run: MemAvailable ${avail} MiB < ${MIN_FREE_MB} MiB -- stopping the job" >&2
    kill -TERM -- "-$pid" 2>/dev/null; sleep 5; kill -KILL -- "-$pid" 2>/dev/null
    tripped=1
    break
  fi
  sleep 1
done
wait "$pid"; rc=$?
(( tripped )) && rc=137
(( rc == 137 && ! tripped )) && [[ $mode == cgroup* ]] \
  && echo "pi_run: the job was killed (SIGKILL); most likely it exceeded MemoryMax=$MEM" >&2
echo "pi_run: exit $rc, peak memory $((peak / 1024)) MiB ($mode, guard ${MIN_FREE_MB} MiB free)" >&2
exit $rc