#!/bin/bash

NR_CPUS=$(getconf _NPROCESSORS_ONLN)
NR_CGROUPS=$(( NR_CPUS * 2 ))
TEST_MB=50
TOTAL_MB=$((TEST_MB * NR_CGROUPS))
TMPFS=$(mktemp -d)
ROOT="/sys/fs/cgroup/"
ZRAM_DEV="/mnt/devtmpfs/zram0"

cleanup() {
  umount $TMPFS
  rm -rf $TMPFS
  for i in $(seq $NR_CGROUPS); do
    cgroup="$ROOT/cg$i"
    rmdir $cgroup
  done
  swapoff $ZRAM_DEV
  echo 1 > "/sys/block/zram0/reset"
}
trap cleanup INT QUIT EXIT

# Setup zram
echo $((TOTAL_MB << 20)) > "/sys/block/zram0/disksize"
mkswap $ZRAM_DEV
swapon $ZRAM_DEV
echo "Setup zram done"

# Create cgroups, set limits
echo "+memory" > "$ROOT/cgroup.subtree_control"
for i in $(seq $NR_CGROUPS); do
  cgroup="$ROOT/cg$i"
  mkdir $cgroup
  echo $(( (TEST_MB << 20) / 4)) > "$cgroup/memory.max"
done
echo "Setup cgroups done"

# Start workers to allocate tmpfs memory
mount -t tmpfs none $TMPFS
for i in $(seq $NR_CGROUPS); do
  cgroup="$ROOT/cg$i"
  f="$TMPFS/tmp$i"
  (echo 0 > "$cgroup/cgroup.procs" &&
    dd if=/dev/zero of=$f bs=1M count=$TEST_MB status=none &&
    cat $f > /dev/null)&
done

# Wait for workers
wait