#!/bin/bash

NR_L1_CGROUPS=50
NR_L2_CGROUPS_PER_L1=5
NR_L2_CGROUPS=$(( NR_L1_CGROUPS * NR_L2_CGROUPS_PER_L1 ))
NR_WORKERS_PER_CGROUP=10
WORKER_MB=10
NR_TOTAL_WORKERS=$(( NR_WORKERS_PER_CGROUP * (1 + NR_L1_CGROUPS + NR_L2_CGROUPS) ))
L1_CGROUP_LIMIT_MB=$(( NR_L2_CGROUPS_PER_L1 * NR_WORKERS_PER_CGROUP * WORKER_MB / 4 ))
TOTAL_MB=$(( NR_TOTAL_WORKERS * WORKER_MB ))
TMPFS=$(mktemp -d)
ROOT="/sys/fs/cgroup/"
ADMIN="/sys/fs/cgroup/admin"
ZRAM_DEV="/mnt/devtmpfs/zram0"

cleanup_cgroup() {
  local -r cg=$1
  local -r procs=$(cat $cg/cgroup.procs)
  if [[ -n $procs ]]; then
    kill -KILL $procs
    wait $procs 2>/dev/null
  fi
  while [[ -n $(cat $cg/cgroup.procs) ]]; do
    sleep 0.1
  done
  rmdir $cg
}

cleanup() {
  for i in $(seq $NR_L1_CGROUPS); do
    l1="$ROOT/parent$i"
    for j in $(seq $NR_L2_CGROUPS_PER_L1); do
      l2="$l1/child$j"
      cleanup_cgroup $l2
    done
    cleanup_cgroup $l1
  done

  cleanup_cgroup $ADMIN

  umount $TMPFS
  rm -rf $TMPFS

  swapoff $ZRAM_DEV
  echo 1 > "/sys/block/zram0/reset"
}
trap cleanup INT QUIT EXIT

run_stats_reader() {
  local -r cg_run=$1
  local -r cg_read=$2

  # read the stats every second until workers are done
  echo 0 > "$cg_run/cgroup.procs"
  while [[ $(ls $TMPFS) ]]; do
    cat "$cg_read/memory.stat" > /dev/null
    sleep 1
  done
}

run_worker() {
  local -r cg=$1
  local -r f=$2

  echo 0 > "$cg/cgroup.procs"
  rm -rf "$f"
  dd if=/dev/zero of="$f" bs=1M count=$WORKER_MB status=none
  cat "$f" > /dev/null
  rm "$f"
}

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

# Mount tmpfs
mount -t tmpfs none $TMPFS

# Create admin cgroup
mkdir $ADMIN

# Create worker cgroups, set limits
echo "+memory" > "$ROOT/cgroup.subtree_control"
for i in $(seq $NR_L1_CGROUPS); do
  l1="$ROOT/parent$i"
  mkdir $l1
  echo "+memory" > "$l1/cgroup.subtree_control"
  for j in $(seq $NR_L2_CGROUPS_PER_L1); do
    l2="$l1/child$j"
    mkdir $l2
  done
  echo $(( L1_CGROUP_LIMIT_MB << 20)) > "$l1/memory.max"
done
echo "Setup $NR_L1_CGROUPS L1 cgroups with limit ${L1_CGROUP_LIMIT_MB}M done"
echo "Each L1 cgroup has $NR_L2_CGROUPS_PER_L1 L2 children"

# Start workers to allocate tmpfs memory
for i in $(seq $NR_L1_CGROUPS); do
  l1="$ROOT/parent$i"
  for j in $(seq $NR_L2_CGROUPS_PER_L1); do
    l2="$l1/child$j"
    for k in $(seq $NR_WORKERS_PER_CGROUP); do
      (run_worker "$l2" "$TMPFS/$i-$j-$k")&
    done
  done
done

# Start stat readers
(run_stats_reader "$ADMIN" "$ROOT")&
for i in $(seq $NR_L1_CGROUPS); do
  l1="$ROOT/parent$i"
  (run_stats_reader "$ADMIN" "$l1")&
  for j in $(seq $NR_L2_CGROUPS_PER_L1); do
    l2="$l1/child$j"
    # Ran stat readers in the admin cgroup as well as the cgroup itself
    (run_stats_reader "$ADMIN" "$l2")&
    (run_stats_reader "$l2" "$l2")&
  done
done

# Wait for workers
echo "Ran $NR_WORKERS_PER_CGROUP workers per L2 cgroup, each allocating ${WORKER_MB}M, waiting.."
wait