ms-runner/starter/cgroup_prepare.c

#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,"%ld\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);
    char* processes_string=NULL;
    asprintf(&processes_string,"%d\n",limits->processes);
    write_file("pids.max",processes_string);
    free(processes_string);
    //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);
}

int check_mem() {
    char cwd[PATH_MAX];
    if(getcwd(cwd,sizeof(cwd))==NULL) die("getcwd error: %m");
    chdir(cgroups_path);
    chdir(cgroup_name);
    int fdmax = open("memory.max",O_RDONLY);
    int fdpeak = open("memory.peak",O_RDONLY);
    char max[101],peak[101];
    read(fdmax,max,100);
    read(fdpeak,peak,100);
    chdir(cwd);
    close(fdmax);
    close(fdpeak);
    if(strcmp(max,peak) == 0) {return -1;}
    return atoi(peak);
}

void remove_cgroup() {
    char cwd[PATH_MAX];
    getcwd(cwd,sizeof(cwd));
    chdir(cgroups_path);
    rmdir(cgroup_name);
    chdir(cwd);
}