#!/bin/bash
#
# This script sets up cgroup2 accounting cgroups for nodes by
# distributing qemu-system-* and gnt-start processes into separate
# groups

## Setup up cgroups2 account, unless done
grep -q "cgroup_root /sys/fs/cgroup tmpfs" /proc/mounts || \
    mount -t tmpfs cgroup_root /sys/fs/cgroup
[ -d /sys/fs/cgroup/cpuset ] || \
    mkdir /sys/fs/cgroup/cpuset
grep -q "cpuset /sys/fs/cgroup/cpuset cgroup2" /proc/mounts || \
    mount -t cgroup2 cpuset /sys/fs/cgroup/cpuset

typeset -A GROUP

## Capture all current groups
for PG in $(ls /sys/fs/cgroup/cpuset/*/cgroup.procs) ; do
    G=${PG%/*}
    G=${G##*/}
    for P in $(cat $PG) ; do
	GROUP[$P]=$G
    done
done

# add pid to group
addtogroup() {
    echo adding $1 to $2 >&2
    GROUP[$1]=$2
    mkdir -p /sys/fs/cgroup/cpuset/$2
    echo $1 > /sys/fs/cgroup/cpuset/$2/cgroup.procs 2>/dev/null
}

qemu_name() {
    < /proc/$1/cmdline xargs -0 -n 1 echo | grep -A1 -e-name | tail -n1 
}

subhost_name() {
    < /proc/$1/cmdline xargs -0 -n 1 echo | tail -n1
}

addtogroup 0 $(hostname -s)

for P in $(pgrep qemu-system) ; do
    addtogroup $P "$(qemu_name $P | sed 's/\..*//')"
done

for P in $(pgrep gnt-start) ; do
    addtogroup $P "$(subhost_name $P | sed 's/\..*//')"
done

for i in {1..100} ; do
    AGAIN=false
    for P in $(cat /sys/fs/cgroup/cpuset/cgroup.procs) ; do
	PP=( $(ps -oppid -p $P | tail -n1) )
	[ -z "$PP" ] && continue
	if [ -z "${GROUP[$PP]}" ] ; then
	    echo "no group for $PP" >&2
	    AGAIN=true
	else
	    addtogroup $P "${GROUP[$PP]}"
	fi
    done
    $AGAIN || break
    echo "again $AGAIN"
done