CpuSet Utility
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

160 lines
3.0 KiB

  1. /**
  2. * @PROJECT CpuSet
  3. * @COPYRIGHT See COPYING in the top level directory
  4. * @FILE cpuset.c
  5. * @PURPOSE Utility for managing QEMU resources
  6. * @DEVELOPERS Rafal Kupiec <belliash@codingworkshop.eu.org>
  7. */
  8. #include <stdio.h>
  9. #include <stdlib.h>
  10. #include <unistd.h>
  11. #include <sys/stat.h>
  12. #define CPUSET(x) "/sys/fs/cgroup/cpuset/"x
  13. #define ACTION_ENABLE 1
  14. #define ACTION_DISABLE 2
  15. char *readLine(FILE *fp, char *buffer) {
  16. int ch;
  17. int i = 0;
  18. size_t buff_len = 0;
  19. buffer = malloc(buff_len + 1);
  20. if(!buffer) {
  21. return NULL;
  22. }
  23. while((ch = fgetc(fp)) != '\n' && ch != EOF) {
  24. buff_len++;
  25. void *tmp = realloc(buffer, buff_len + 1);
  26. if(tmp == NULL) {
  27. free(buffer);
  28. return NULL;
  29. }
  30. buffer = tmp;
  31. buffer[i] = (char) ch;
  32. i++;
  33. }
  34. buffer[i] = '\0';
  35. if(ch == EOF && (i == 0 || ferror(fp))) {
  36. free(buffer);
  37. return NULL;
  38. }
  39. return buffer;
  40. }
  41. int moveTasks(char *src, char *dst) {
  42. FILE *fd1, *fd2;
  43. char *pid;
  44. fd1 = fopen(src, "rb");
  45. if(fd1 == NULL) {
  46. return 255;
  47. }
  48. while((pid = readLine(fd1, 0)) != NULL) {
  49. if(strtol(pid, (char **)NULL, 10) != getppid()) {
  50. fd2 = fopen(dst, "wb");
  51. if(fd2 == NULL) {
  52. return 255;
  53. }
  54. fprintf(fd2, "%s\n", pid);
  55. fclose(fd2);
  56. free(pid);
  57. }
  58. }
  59. fclose(fd1);
  60. return 0;
  61. }
  62. int setGroup(char *dest, char *data) {
  63. FILE *fd;
  64. fd = fopen(dest, "wb");
  65. if(fd == NULL) {
  66. return 255;
  67. }
  68. fprintf(fd, "%s", data);
  69. fclose(fd);
  70. return 0;
  71. }
  72. int disableGroup() {
  73. struct stat sb = {0};
  74. printf("Disabling CGROUPS ... ");
  75. if(stat(CPUSET("host"), &sb) == 0) {
  76. if(moveTasks(CPUSET("host/tasks"), CPUSET("tasks")) != 0) {
  77. return 255;
  78. }
  79. rmdir(CPUSET("host"));
  80. printf("OK!\n");
  81. } else {
  82. printf("SKIPPED!\n");
  83. }
  84. return 0;
  85. }
  86. int enableGroup(char *cpu, char *mem) {
  87. struct stat sb = {0};
  88. printf("Enabling CGROUPS ... ");
  89. if(stat(CPUSET("host"), &sb) == -1) {
  90. mkdir(CPUSET("host"), 0755);
  91. }
  92. if(setGroup(CPUSET("host/cpuset.cpus"), cpu) != 0) {
  93. printf("FAIL!\n");
  94. return 255;
  95. }
  96. if(setGroup(CPUSET("host/cpuset.mems"), mem) != 0) {
  97. printf("FAIL!\n");
  98. return 255;
  99. }
  100. if(moveTasks(CPUSET("tasks"), CPUSET("host/tasks")) != 0) {
  101. printf("FAIL!\n");
  102. return 255;
  103. }
  104. printf("OK!\n");
  105. return 0;
  106. }
  107. int main(int argc, char *argv[]) {
  108. int opt;
  109. int action = 0;
  110. char *cpu = {0};
  111. char *mem = {0};
  112. while((opt = getopt(argc, argv, ":edc:m:")) != -1) {
  113. switch(opt) {
  114. case 'e':
  115. action = ACTION_ENABLE;
  116. break;
  117. case 'd':
  118. action = ACTION_DISABLE;
  119. break;
  120. case 'c':
  121. cpu = optarg;
  122. break;
  123. case 'm':
  124. mem = optarg;
  125. break;
  126. case ':':
  127. printf("Option needs a value!\n");
  128. return 1;
  129. default:
  130. printf("Invalid option: -%c!\n", optopt);
  131. return 1;
  132. }
  133. }
  134. if(geteuid() != 0) {
  135. printf("This program needs to be launched as root!\n");
  136. return 1;
  137. }
  138. if(action == ACTION_ENABLE) {
  139. if(!cpu || !mem) {
  140. printf("No CPU or memory pinning specified!\n");
  141. return 1;
  142. }
  143. return enableGroup(cpu, mem);
  144. } else if(action == ACTION_DISABLE) {
  145. return disableGroup();
  146. } else {
  147. printf("No action specified!\n");
  148. }
  149. return 0;
  150. }