Add prototype impl for cgroupv2

Author: meisterT

judge/runguard.cc

@@ -45,6 +45,7 @@
 #include <fcntl.h>
 #include <csignal>
 #include <cstdlib>
+#include <mntent.h>
 #include <unistd.h>
 #include <cstring>
 #include <cstdarg>
@@ -138,6 +139,8 @@ int show_help;
 int show_version;
 pid_t runpipe_pid = -1;
 
+bool is_cgroup_v2 = false;
+
 double walltimelimit[2], cputimelimit[2]; /* in seconds, soft and hard limits */
 int walllimit_reached, cpulimit_reached; /* 1=soft, 2=hard, 3=both limits reached */
 int64_t memsize;
@@ -433,7 +436,11 @@ void output_exit_time(int exitcode, double cpudiff)
 void check_remaining_procs()
 {
 	char path[1024];
-	snprintf(path, 1023, "/sys/fs/cgroup/cpuacct%scgroup.procs", cgroupname);
+    if (is_cgroup_v2) {
+        snprintf(path, 1023, "/sys/fs/cgroup/%scgroup.procs", cgroupname);
+    } else {
+        snprintf(path, 1023, "/sys/fs/cgroup/cpuacct%scgroup.procs", cgroupname);
+    }
 
 	FILE *file = fopen(path, "r");
 	if (file == nullptr) {
@@ -447,7 +454,7 @@ void check_remaining_procs()
 	if (fclose(file) != 0) error(errno, "closing file `%s'", path);
 }
 
-void output_cgroup_stats(double *cputime)
+void output_cgroup_stats_v1(double *cputime)
 {
 	struct cgroup *cg;
 	if ( (cg = cgroup_new_cgroup(cgroupname))==nullptr ) error(0,"cgroup_new_cgroup");
@@ -474,6 +481,45 @@ void output_cgroup_stats(double *cputime)
 	cgroup_free(&cg);
 }
 
+void output_cgroup_stats_v2(double *cputime)
+{
+    struct cgroup *cg;
+    if ( (cg = cgroup_new_cgroup(cgroupname))==NULL ) error(0,"cgroup_new_cgroup");
+
+    int ret;
+    if ((ret = cgroup_get_cgroup(cg)) != 0) error(ret,"get cgroup information");
+
+    struct cgroup_controller *cg_controller = cgroup_get_controller(cg, "memory");
+    int64_t max_usage = 0;
+    ret = cgroup_get_value_int64(cg_controller, "memory.peak", &max_usage);
+    if ( ret == ECGROUPVALUENOTEXIST ) {
+        write_meta("internal-warning", "Kernel too old and does not support memory.peak");
+    } else if ( ret!=0 ) {
+        error(ret,"get cgroup value memory.peak");
+    }
+
+    // There is no need to check swap usage, as we limit it to 0.
+    verbose("total memory used: %" PRId64 " kB", max_usage/1024);
+    write_meta("memory-bytes","%" PRId64, max_usage);
+
+    struct cgroup_stat stat;
+    void *handle;
+    ret = cgroup_read_stats_begin("cpu", cgroupname, &handle, &stat);
+    while (ret == 0) {
+        if (strcmp(stat.name, "usage_usec") == 0) {
+            long long usec = strtoll(stat.value, NULL, 10);
+            *cputime = usec / 1e6;
+            break;
+        }
+        ret = cgroup_read_stats_next(&handle, &stat);
+    }
+    if ( ret!=ECGEOF ) error(ret,"get cgroup value cpu.stat");
+    cgroup_read_stats_end(&handle);
+
+    cgroup_free(&cg);
+
+}
+
 /* Temporary shorthand define for error handling. */
 #define cgroup_add_value(type,name,value) \
 	ret = cgroup_add_value_ ## type(cg_controller, name, value); \
@@ -492,9 +538,21 @@ void cgroup_create()
 		error(0,"cgroup_add_controller memory");
 	}
 
-	int ret;
-	cgroup_add_value(int64, "memory.limit_in_bytes", memsize);
-	cgroup_add_value(int64, "memory.memsw.limit_in_bytes", memsize);
+    int ret;
+    if (is_cgroup_v2) {
+        // TODO: do we want to set cpu.weight here as well?
+        if (memsize != RLIM_INFINITY) {
+            cgroup_add_value(int64, "memory.max", memsize);
+            // TODO: Is this the behavior change that JohnB mentioned?
+            cgroup_add_value(int64, "memory.swap.max", 0);
+        } else {
+            cgroup_add_value(string, "memory.max", "max");
+            cgroup_add_value(string, "memory.swap.max", "max");
+        }
+    } else {
+        cgroup_add_value(int64, "memory.limit_in_bytes", memsize);
+        cgroup_add_value(int64, "memory.memsw.limit_in_bytes", memsize);
+    }
 
 	/* Set up cpu restrictions; we pin the task to a specific set of
 	   cpus. We also give it exclusive access to those cores, and set
@@ -512,9 +570,14 @@ void cgroup_create()
 		verbose("cpuset undefined");
 	}
 
-	if ( (cg_controller = cgroup_add_controller(cg, "cpuacct"))==nullptr ) {
-		error(0,"cgroup_add_controller cpuacct");
+	if ( (cg_controller = cgroup_add_controller(cg, "cpu"))==nullptr ) {
+		error(0,"cgroup_add_controller cpu");
 	}
+    if (!is_cgroup_v2) {
+        if ((cg_controller = cgroup_add_controller(cg, "cpuacct")) == nullptr) {
+            error(0, "cgroup_add_controller cpuacct");
+        }
+    }
 
 	/* Perform the actual creation of the cgroup */
 	if ( (ret = cgroup_create_cgroup(cg, 1))!=0 ) error(ret,"creating cgroup");
@@ -542,16 +605,28 @@ void cgroup_attach()
 
 void cgroup_kill()
 {
-	void *handle = nullptr;
-	pid_t pid;
-
-	/* kill any remaining tasks, and wait for them to be gone */
-	while(1) {
-		int ret = cgroup_get_task_begin(cgroupname, "memory", &handle, &pid);
-		cgroup_get_task_end(&handle);
-		if (ret == ECGEOF) break;
-		kill(pid, SIGKILL);
-	}
+    /* kill any remaining tasks, and wait for them to be gone */
+    if (is_cgroup_v2) {
+        int size;
+        do {
+            pid_t* pids;
+            int ret = cgroup_get_procs(cgroupname, NULL, &pids, &size);
+            if (ret != 0) error(ret, "cgroup_get_procs");
+            for(int i = 0; i < size; i++) {
+                kill(pids[i], SIGKILL);
+            }
+            free(pids);
+        } while (size > 0);
+    } else {
+        while(1) {
+            void *handle = nullptr;
+            pid_t pid;
+            int ret = cgroup_get_task_begin(cgroupname, "memory", &handle, &pid);
+            cgroup_get_task_end(&handle);
+            if (ret == ECGEOF) break;
+            kill(pid, SIGKILL);
+        }
+    }
 }
 
 void cgroup_delete()
@@ -560,7 +635,10 @@ void cgroup_delete()
 	cg = cgroup_new_cgroup(cgroupname);
 	if (!cg) error(0,"cgroup_new_cgroup");
 
-	if ( cgroup_add_controller(cg, "cpuacct")==nullptr ) error(0,"cgroup_add_controller cpuacct");
+    if (cgroup_add_controller(cg, "cpu") == nullptr) error(0, "cgroup_add_controller cpu");
+    if (!is_cgroup_v2) {
+        if (cgroup_add_controller(cg, "cpuacct") == nullptr) error(0, "cgroup_add_controller cpuacct");
+    }
 	if ( cgroup_add_controller(cg, "memory")==nullptr ) error(0,"cgroup_add_controller memory");
 
 	if ( cpuset!=nullptr && strlen(cpuset)>0 ) {
@@ -783,7 +861,13 @@ void setrestrictions()
 	}
 
 	/* Put the child process in the cgroup */
-	cgroup_attach();
+    if (is_cgroup_v2) {
+       if (cgroup_change_cgroup_path(cgroupname, getpid(), NULL) != 0) {
+           error(0, "Failed to move the process to the cgroup");
+       }
+    } else {
+        cgroup_attach();
+    }
 
 	/* Run the command in a separate process group so that the command
 	   and all its children can be killed off with one signal. */
@@ -929,6 +1013,29 @@ void pump_pipes(fd_set* readfds, size_t data_read[], size_t data_passed[])
 
 }
 
+bool cgroup_is_v2() {
+    bool ret = false;
+    FILE *fp = setmntent("/proc/mounts", "r");
+    if (!fp) {
+        perror("Error opening /proc/mounts");
+        return false;
+    }
+
+    struct mntent *entry;
+    while ((entry = getmntent(fp)) != nullptr) {
+        if (strcmp(entry->mnt_dir, "/sys/fs/cgroup") == 0) {
+            if (strcmp(entry->mnt_type, "cgroup2") == 0) {
+                ret = true;
+            }
+            break;
+        }
+    }
+
+    endmntent(fp);
+
+    return ret;
+}
+
 int main(int argc, char **argv)
 {
 	int   ret;
@@ -1103,6 +1210,8 @@ int main(int argc, char **argv)
 	cmdname = argv[optind];
 	cmdargs = argv+optind;
 
+    is_cgroup_v2 = cgroup_is_v2();
+
 	if ( outputmeta && (metafile = fopen(metafilename,"w"))==nullptr ) {
 		error(errno,"cannot open `%s'",metafilename);
 	}
@@ -1162,6 +1271,7 @@ int main(int argc, char **argv)
 			}
 		}
 	}
+
 	/* Make libcgroup ready for use */
 	ret = cgroup_init();
 	if ( ret!=0 ) {
@@ -1175,8 +1285,8 @@ int main(int argc, char **argv)
 	} else {
 		str[0] = 0;
 	}
-	snprintf(cgroupname, 255, "/domjudge/dj_cgroup_%d_%.16s_%d.%06d/",
-	         getpid(), str, (int)progstarttime.tv_sec, (int)progstarttime.tv_usec);
+    snprintf(cgroupname, 255, "/domjudge/dj_cgroup_%d_%.16s_%d.%06d/",
+             getpid(), str, (int) progstarttime.tv_sec, (int) progstarttime.tv_usec);
 
 	cgroup_create();
 
@@ -1422,7 +1532,11 @@ int main(int argc, char **argv)
 		check_remaining_procs();
 
 		double cputime = -1;
-		output_cgroup_stats(&cputime);
+        if (is_cgroup_v2) {
+            output_cgroup_stats_v2(&cputime);
+        } else {
+            output_cgroup_stats_v1(&cputime);
+        }
 		cgroup_kill();
 		cgroup_delete();