#!/bin/bash
#
# Mounted with:
#   mount -t cgroup2 -o memory_recursiveprot none $ROOT
# Or remounted with:
#   mount -o remount,memory_recursiveprot $ROOT

# Default root path
: ${ROOT:="/sys/fs/cgroup"}
A="$ROOT/A"
B="$A/B"

# Default to current directory
: ${TEST_DIR:=.}
SWAPFILE=
TMPFS=
PIPE=

setup() {
  set -e

  # Setup 100M swapfile
  SWAPFILE=${TEST_DIR}/test_swapfile
  dd if=/dev/zero of=$SWAPFILE bs=1M count=100 status=none
  chmod 600 $SWAPFILE
  mkswap $SWAPFILE > /dev/null
  swapon $SWAPFILE

  # Setup tmpfs
  TMPFS=${TEST_DIR}/test_tmpfs
  mkdir $TMPFS
  mount -t tmpfs tmpfs $TMPFS

  # Setup pipe
  PIPE=${TEST_DIR}/test_pipe
  mkfifo $PIPE

  # Setup root
  enable_memcg_subtree $ROOT

  # Setup A
  mkdir $A
  enable_memcg_subtree $A
  echo 50M > "$A/memory.min"

  # Setup B
  mkdir $B
  echo 10M > "$B/memory.min"
  echo 40M > "$B/memory.high"
}

cleanup() {
  set +e

  # Kill any procs in B first, then remove it.
  procs=$(cat $B/cgroup.procs)
  if [[ -n $procs ]]; then
    kill -KILL $procs
    wait $procs 2>/dev/null
  fi
  while [[ -n $(cat $B/cgroup.procs) ]]; do
    sleep 0.1
  done
  rmdir $B

  # Remove A
  rmdir $A

  # Cleanup everything else
  if [[ -n $TMPFS ]]; then
    umount $TMPFS
    rmdir $TMPFS
  fi
  swapoff $SWAPFILE
  rm $SWAPFILE
  rm $PIPE
}

fail() {
  echo "[FAIL] $*"
  exit 1
}

pass() {
  echo "[PASS] $*"
}

allocate_tmpfs() {
  echo 0 > $1/cgroup.procs
  dd if=/dev/zero bs=1M count=$2 >> $TMPFS/file status=none
  echo "tmpfs" > $PIPE
  sleep 1000
}

wait_for_pipe() {
  if read -t 10 line < $PIPE; then
    if [[ $line == "tmpfs" ]]; then
      return
    else
      fail "wrong input ($line) received on pipe, expected (tmpfs)"
    fi
  fi
  fail "pipe timoeut"
}

memcg_usage() {
  cat $1/memory.current
}

enable_memcg_subtree() {
  echo "+memory" > "$1/cgroup.subtree_control"
}

trap cleanup EXIT

echo "Reference case"
# A: memory.min = 50M
# A/B: memory.min = 10M, memory.high = 40M
# Allocate 35 M in B, nothing happens
# Allocate 10 M more in B, memory.high pushes us back below 40 M
setup

(allocate_tmpfs $B 35)&
wait_for_pipe

usage_mb=$(( $(memcg_usage $B) / (1024 * 1024) ))
echo "B's usage after initial allocation: $usage_mb M"

(allocate_tmpfs $B 10)&
wait_for_pipe

usage_mb=$(( $(memcg_usage $B) / (1024 * 1024) ))
echo "B's usage after allocation beyond high: $usage_mb M"

if [[ $usage_mb -le 40 ]]; then
  pass "memory.high enforced effectively"
else
  fail "memory.high not enforced"
fi

cleanup

echo "Corner case"
# Corner case
# Same setup as above
# Allocate 35 M in B, nothing happens
# Invoke global reclaim, all of B's memory should be protected by A's memory.min
# Allocate 10 M more in B, memory high should push us back below 40 M
setup

(allocate_tmpfs $B 35)&
wait_for_pipe

usage_mb=$(( $(memcg_usage $B) / (1024 * 1024) ))
echo "B's usage after initial allocation: $usage_mb M"

# Simulate global reclaim
set +e
echo 10M > $ROOT/memory.reclaim
set -e

usage_mb=$(( $(memcg_usage $B) / (1024 * 1024) ))
echo "B's usage after global reclaim: $usage_mb M"

# B should be fully protected from global reclaim by A's min
usage_mb_after=$(( $(memcg_usage $B) / (1024 * 1024) ))
if [[ $usage_mb_after -ne $usage_mb ]]; then
  fail "A's memory.min did not protect B from global reclaim"
else
  pass "A's memory.min protected B from global reclaim"
fi

(allocate_tmpfs $B 10)&
wait_for_pipe

# B's memory.high should be enforced and it should remain below 40M
usage_mb=$(( $(memcg_usage $B) / (1024 * 1024) ))
echo "B's usage after global reclaim: $usage_mb M"

if [[ $usage_mb -le 40 ]]; then
  pass "memory.high enforced effectively"
else
  fail "memory.high not enforced"
fi