CPULimiteHost.popen(): set cgroup and (optionally) RT priority

mnexec.c

@@ -7,6 +7,8 @@
  *  - detaching from a controlling tty using setsid
  *  - running in a network namespace
  *  - printing out the pid of a process so we can identify it later
+ *  - attaching to a namespace and cgroup
+ *  - setting RT scheduling
  *
  * Partially based on public domain setsid(1)
 */
@@ -17,16 +19,21 @@
 #include <limits.h>
 #include <syscall.h>
 #include <fcntl.h>
+#include <stdlib.h>
+#include <limits.h>
+#include <sched.h>
 
 void usage(char *name) 
 {
     printf("Execution utility for Mininet.\n"
-           "usage: %s [-cdnp]\n"
+           "usage: %s [-cdnp] [-a pid] [-g group] [-r rtprio] cmd args...\n"
            "-c: close all file descriptors except stdin/out/error\n"
            "-d: detach from tty by calling setsid()\n"
            "-n: run in new network namespace\n"
            "-p: print ^A + pid\n"
-           "-a pid: attach to pid's network namespace\n", 
+           "-a pid: attach to pid's network namespace\n"
+           "-g group: add to cgroup\n"
+           "-r rtprio: run with SCHED_RR (usually requires -g)\n",
            name);
 }
 
@@ -36,6 +43,47 @@ int setns(int fd, int nstype)
 	return syscall(308, fd, nstype);
 }
 
+/* Validate alphanumeric path foo1/bar2/baz */
+void validate(char *path)
+{
+    char *s;
+    for (s=path; *s; s++) {
+        if (!isalnum(*s) && *s != '/') {
+            fprintf(stderr, "invalid path: %s\n", path);
+            exit(1);
+        }
+    }
+}
+
+/* Add our pid to cgroup */
+int cgroup(char *gname)
+{
+    static char path[PATH_MAX];
+    static char *groups[] = {
+        "cpu", "cpuacct", "cpuset", NULL
+    };
+    char **gptr;
+    pid_t pid = getpid();
+    int count = 0;
+    validate(gname);
+    for (gptr = groups; *gptr; gptr++) {
+        FILE *f;
+        snprintf(path, PATH_MAX, "/sys/fs/cgroup/%s/%s/tasks",
+                 *gptr, gname);
+        f = fopen(path, "w");
+        if (f) {
+            count++;
+            fprintf(f, "%d\n", pid);
+            fclose(f);
+        }
+    }
+    if (!count) {
+        fprintf(stderr, "cgroup: could not add to cgroup %s\n",
+            gname);
+        exit(1);
+    }
+}
+
 int main(int argc, char *argv[])
 {
     char c;
@@ -43,8 +91,8 @@ int main(int argc, char *argv[])
     char path[PATH_MAX];
     int nsid;
     int pid;
-
-    while ((c = getopt(argc, argv, "+cdnpa:")) != -1)
+    static struct sched_param sp;
+    while ((c = getopt(argc, argv, "+cdnpa:g:r:")) != -1)
         switch(c) {
         case 'c':
             /* close file descriptors except stdin/out/error */
@@ -92,16 +140,28 @@ int main(int argc, char *argv[])
                 return 1;
             }
             break;
+        case 'g':
+            /* Attach to cgroup */
+            cgroup(optarg);
+            break;
+        case 'r':
+            /* Set RT scheduling priority */
+            sp.sched_priority = atoi(optarg);
+            if (sched_setscheduler(getpid(), SCHED_RR, &sp) < 0) {
+                perror("sched_setscheduler");
+                return 1;
+            }
+            break;
         default:
             usage(argv[0]);
             break;
         }
 
     if (optind < argc) {
-		execvp(argv[optind], &argv[optind]);
-		perror(argv[optind]);
-		return 1;
-	}
+        execvp(argv[optind], &argv[optind]);
+        perror(argv[optind]);
+        return 1;
+    }
     
     usage(argv[0]);
 }