associate SSH processes with the appropriate cgroup

[util-vserver-pl.git] / src / vsh.c

diff --git a/src/vsh.c b/src/vsh.c
index 1ab413e..ef6dfa8 100644 (file)
--- a/src/vsh.c
+++ b/src/vsh.c
@@ -45,6 +45,41 @@ char **extend_argv(int argc, char **argv, int num_extra_args) {
     return argv2;
 }
 
+void associate_vserver_cgroup(char *slice_name, int pid)
+{
+    char cgroup_dir[4096], fn[4096];
+    struct stat st;
+    int result;
+    FILE *f;
+
+    sprintf(cgroup_dir, "/dev/cgroup/%s", slice_name);
+
+    result = stat(cgroup_dir, &st);
+    if (result != 0) {
+        // doesn't exist
+        return;
+    }
+
+    if (!S_ISDIR(st.st_mode)) {
+        // not a directory
+        return;
+    }
+
+    /* Write the pid to the cgroup tasks file, so the SSH process is associated
+     * with the correct cgroup
+     */
+
+    sprintf(fn, "%s/tasks", cgroup_dir);
+    f = fopen(fn, "wt");
+    if (f==NULL) {
+        // failed to open
+        return;
+    }
+
+    fprintf(f, "%d\n", pid);
+    fclose(f);
+}
+
 #define NUM_VSERVER_EXEC_ARGS 5
 
 int main(int argc, char **argv, char **envp)
@@ -61,6 +96,8 @@ int main(int argc, char **argv, char **envp)
         fprintf(stderr,"Could not look up slice name\n");
         goto out_exception;
     }
+
+    associate_vserver_cgroup(slice_name, getpid());
     
     argv2 = extend_argv(argc, argv, NUM_VSERVER_EXEC_ARGS);
     if (!argv2) goto out_exception;