160 rader
3.0 KiB
C
160 rader
3.0 KiB
C
/**
|
|
* @PROJECT CpuSet
|
|
* @COPYRIGHT See COPYING in the top level directory
|
|
* @FILE cpuset.c
|
|
* @PURPOSE Utility for managing QEMU resources
|
|
* @DEVELOPERS Rafal Kupiec <belliash@codingworkshop.eu.org>
|
|
*/
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <unistd.h>
|
|
#include <sys/stat.h>
|
|
|
|
#define CPUSET(x) "/sys/fs/cgroup/cpuset/"x
|
|
|
|
#define ACTION_ENABLE 1
|
|
#define ACTION_DISABLE 2
|
|
|
|
char *readLine(FILE *fp, char *buffer) {
|
|
int ch;
|
|
int i = 0;
|
|
size_t buff_len = 0;
|
|
buffer = malloc(buff_len + 1);
|
|
if(!buffer) {
|
|
return NULL;
|
|
}
|
|
while((ch = fgetc(fp)) != '\n' && ch != EOF) {
|
|
buff_len++;
|
|
void *tmp = realloc(buffer, buff_len + 1);
|
|
if(tmp == NULL) {
|
|
free(buffer);
|
|
return NULL;
|
|
}
|
|
buffer = tmp;
|
|
buffer[i] = (char) ch;
|
|
i++;
|
|
}
|
|
buffer[i] = '\0';
|
|
if(ch == EOF && (i == 0 || ferror(fp))) {
|
|
free(buffer);
|
|
return NULL;
|
|
}
|
|
return buffer;
|
|
}
|
|
|
|
int moveTasks(char *src, char *dst) {
|
|
FILE *fd1, *fd2;
|
|
char *pid;
|
|
fd1 = fopen(src, "rb");
|
|
if(fd1 == NULL) {
|
|
return 255;
|
|
}
|
|
while((pid = readLine(fd1, 0)) != NULL) {
|
|
if(strtol(pid, (char **)NULL, 10) != getppid()) {
|
|
fd2 = fopen(dst, "wb");
|
|
if(fd2 == NULL) {
|
|
return 255;
|
|
}
|
|
fprintf(fd2, "%s\n", pid);
|
|
fclose(fd2);
|
|
free(pid);
|
|
}
|
|
}
|
|
fclose(fd1);
|
|
return 0;
|
|
}
|
|
|
|
int setGroup(char *dest, char *data) {
|
|
FILE *fd;
|
|
fd = fopen(dest, "wb");
|
|
if(fd == NULL) {
|
|
return 255;
|
|
}
|
|
fprintf(fd, "%s", data);
|
|
fclose(fd);
|
|
return 0;
|
|
}
|
|
|
|
int disableGroup() {
|
|
struct stat sb = {0};
|
|
printf("Disabling CGROUPS ... ");
|
|
if(stat(CPUSET("host"), &sb) == 0) {
|
|
if(moveTasks(CPUSET("host/tasks"), CPUSET("tasks")) != 0) {
|
|
return 255;
|
|
}
|
|
rmdir(CPUSET("host"));
|
|
printf("OK!\n");
|
|
} else {
|
|
printf("SKIPPED!\n");
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int enableGroup(char *cpu, char *mem) {
|
|
struct stat sb = {0};
|
|
printf("Enabling CGROUPS ... ");
|
|
if(stat(CPUSET("host"), &sb) == -1) {
|
|
mkdir(CPUSET("host"), 0755);
|
|
}
|
|
if(setGroup(CPUSET("host/cpuset.cpus"), cpu) != 0) {
|
|
printf("FAIL!\n");
|
|
return 255;
|
|
}
|
|
if(setGroup(CPUSET("host/cpuset.mems"), mem) != 0) {
|
|
printf("FAIL!\n");
|
|
return 255;
|
|
}
|
|
if(moveTasks(CPUSET("tasks"), CPUSET("host/tasks")) != 0) {
|
|
printf("FAIL!\n");
|
|
return 255;
|
|
}
|
|
printf("OK!\n");
|
|
return 0;
|
|
}
|
|
|
|
int main(int argc, char *argv[]) {
|
|
int opt;
|
|
int action = 0;
|
|
char *cpu = {0};
|
|
char *mem = {0};
|
|
while((opt = getopt(argc, argv, ":edc:m:")) != -1) {
|
|
switch(opt) {
|
|
case 'e':
|
|
action = ACTION_ENABLE;
|
|
break;
|
|
case 'd':
|
|
action = ACTION_DISABLE;
|
|
break;
|
|
case 'c':
|
|
cpu = optarg;
|
|
break;
|
|
case 'm':
|
|
mem = optarg;
|
|
break;
|
|
case ':':
|
|
printf("Option needs a value!\n");
|
|
return 1;
|
|
default:
|
|
printf("Invalid option: -%c!\n", optopt);
|
|
return 1;
|
|
}
|
|
}
|
|
if(geteuid() != 0) {
|
|
printf("This program needs to be launched as root!\n");
|
|
return 1;
|
|
}
|
|
if(action == ACTION_ENABLE) {
|
|
if(!cpu || !mem) {
|
|
printf("No CPU or memory pinning specified!\n");
|
|
return 1;
|
|
}
|
|
return enableGroup(cpu, mem);
|
|
} else if(action == ACTION_DISABLE) {
|
|
return disableGroup();
|
|
} else {
|
|
printf("No action specified!\n");
|
|
}
|
|
return 0;
|
|
}
|