#!/usr/bin/perl -w

use strict;

my $cg_mnt    	= "/cgroup";
my $cg_system 	= "memory";
my $node 	= "node";
my $uid 	= "1003";
my $gid 	= "1003";
my $job_id 	= "12345";
my $num_cgroups = 30;
my $job_mem 	= 536870912;
my $step_mem 	= 31457280;
my @cgroups 	= ();
my $sleep_time  = 30;

die "Not running as root" if ($<);

print "I am pid $$\n";

init_ns();
fork_children();
print "Done\n";

sub fork_children
{
	my $batch_cg = create_cg("batch");
	for (1..$num_cgroups) {
		my $sleep = $sleep_time + 15 - $_;
		my $srun_pid = fork();
		if (!$srun_pid) {
			print "srun_$_: $$\n";
			# srun equiv
			my $step_cg = create_cg($_);
			my $exec_pid = fork();
			if (!$exec_pid) {
				print "\tsleep_$_: $$\n";
				exec("sleep $sleep");
				exit(0);
			}
			add_pid($step_cg->{step}, $exec_pid);
			wait();
			destroy_cg($step_cg);
			print "Child $_ finished\n";
			exit(0);
		}
		add_pid($batch_cg->{step}, $srun_pid);
	}

	while (wait() != -1) {}
	destroy_cg($batch_cg);
	print "All children finished\n";
}

sub init_ns
{
	`mkdir -p $cg_mnt/$cg_system`;
	`mount -t cgroup -o $cg_system $cg_system $cg_mnt/$cg_system`;
	`echo 1 > $cg_mnt/$cg_system/memory.use_hierarchy`;
}

sub create_cg
{
	my $step_id = shift;

	#create root cgroup -- /cgroups/memory/node
	my $cg_root = "$cg_mnt/$cg_system/$node";
	cgroup_init($cg_root, 0, $<, $();

	#create user cgroup -- /cgroups/memory/node/uid_#
	my $cg_user = "$cg_root/uid_$uid";
	cgroup_init($cg_user, 1, $<, $();
	`echo 1 > $cg_user/memory.use_hierarchy`;

	#create job cgroup -- /cgroups/memory/node/uid_#/job_#
	my $cg_job = "$cg_user/job_$job_id";
	cgroup_init($cg_job, 1, $<, $();
	mem_init($cg_job, $job_mem);

	#create job cgroup -- /cgroups/memory/node/uid_#/job_#/step_#
	my $cg_step = "$cg_job/step_$step_id";
	cgroup_init($cg_step, 0, $uid, $gid);
	mem_init($cg_step, $step_mem);


	my %cgroup = (step => $cg_step,
		      job  => $cg_job,
		      user => $cg_user);

	return \%cgroup;
}

sub destroy_cg
{
	my $cgroup = shift;

	rmdir $cgroup->{step} or print "$cgroup->{step}: $!\n";
	rmdir $cgroup->{job} or print "$cgroup->{job}: $!\n";
	rmdir $cgroup->{user} or print "$cgroup->{user}: $!\n";
}

sub cgroup_init
{
	my $cg_path = shift;
	my $notify  = shift;
	my $uid     = shift;
	my $gid     = shift;
	#create root cgroup -- /cgroups/memory/node
	mkdir $cg_path unless (-d $cg_path);

	chown $uid, $gid, $cg_path or die "chown: $!";

	if ($notify) {
		`echo 1 > $cg_path/notify_on_release`;
	}
}

sub mem_init
{
	my $cg_path = shift;
	my $limit = shift;
	`echo 1 > $cg_path/memory.use_hierarchy`;
	`echo $limit > $cg_path/memory.limit_in_bytes`;
	`echo $limit > $cg_path/memory.kmem.limit_in_bytes`;
}

sub add_pid
{
	my $cgroup = shift;
	my $pid    = shift;
	`echo $pid > $cgroup/cgroup.procs`;
}
