#include<sys/prctl.h>
#include<sys/mount.h>
#include<sys/stat.h>
#include<unistd.h>
#include<linux/limits.h>

#define CGROUP_NAME_SIZE 20
char cgroup_name[CGROUP_NAME_SIZE+1];
const char* cgroups_path="/sys/fs/cgroup/";

void prepare_cgroup(struct limits* limits) {
    for(int i=0;i<CGROUP_NAME_SIZE;i++) {
        cgroup_name[i]='a'+rand()%26;
    }
    char cwd[PATH_MAX];
    if(getcwd(cwd,sizeof(cwd))==NULL) die("getcwd error: %m");
    chdir(cgroups_path);
    mkdir(cgroup_name,0755);
    chdir(cgroup_name);
    char* memory_string=NULL;
    asprintf(&memory_string,"%d\n",limits->memory);
    write_file("memory.max",memory_string);
    free(memory_string);
    char* cpus_string=NULL;
    asprintf(&cpus_string,"%d\n",limits->core);
    write_file("cpuset.cpus",cpus_string);
    free(cpus_string);
    write_file("pids.max","1\n");
    //write_file("cpuset.cpus","3\n");
    chdir(cwd);
}

void add_to_cgroup(int pid) {
    char cwd[PATH_MAX];
    if(getcwd(cwd,sizeof(cwd))==NULL) die("getcwd error: %m");
    chdir(cgroups_path);
    chdir(cgroup_name);
    char* pidstr=NULL;
    asprintf(&pidstr,"%d\n",pid);
    write_file("cgroup.procs",pidstr);
    free(pidstr);
    chdir(cwd);
}

void remove_cgroup() {
    char cwd[PATH_MAX];
    if(getcwd(cwd,sizeof(cwd))==NULL)   die("getcwd error: %m");
    if(chdir(cgroups_path))             die("chdir error: %m");
    if(rmdir(cgroup_name))              die("rmdir error: %m");
    if(chdir(cwd))                      die("chdir error: %m");
}