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.

148 lines
2.7 KiB

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