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.

243 lines
7.3 KiB

  1. /*
  2. * Copyright (c) 2015 Franco (nextime) Lanza <nextime@nexlab.it>
  3. *
  4. * pinthread [https://git.devuan.org/packages-base/pinthread]
  5. *
  6. * pinthread is free software: you can redistribute it and/or modify
  7. * it under the terms of the GNU General Public License as published by
  8. * the Free Software Foundation, either version 3 of the License, or
  9. * (at your option) any later version.
  10. *
  11. * This program is distributed in the hope that it will be useful,
  12. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. * GNU General Public License for more details.
  15. *
  16. * You should have received a copy of the GNU General Public License
  17. * along with this program. If not, see <http://www.gnu.org/licenses/>.
  18. *
  19. *************************************************************************************
  20. *
  21. * compile with :
  22. * $ gcc -Wall -D_GNU_SOURCE -fpic -shared -o pinthread.so pinthread.c -ldl -lpthread
  23. *
  24. * and test it with :
  25. * LD_PRELOAD=/path/to/pinthread.so something
  26. *
  27. * Environment variables:
  28. * - PINTHREAD_CORE: set to the core number you want to pin to.
  29. * if omitted, it defaults to the result of
  30. * sched_getcpu()
  31. *
  32. * - PINTHREAD_PNAMES: accept a space separated list of strings,
  33. * if any string match the process name or
  34. * PINTHREAD_PNAMES is empty/omitted, it will
  35. * pin all thread according to PINTHREAD_CORE,
  36. * otherwise it will just call the "real"
  37. * pthread_create without any pinning.
  38. *
  39. * - PINTHREAD_DEBUG: if set, enable debug messages on stderr.
  40. */
  41. //#include <sys/types.h>
  42. #include <unistd.h>
  43. #include <stdlib.h>
  44. #include <stdio.h>
  45. #include <pthread.h>
  46. #include <dlfcn.h>
  47. #include <sched.h> //cpu_set_t , CPU_SET
  48. #include <libgen.h>
  49. #include <stdbool.h>
  50. #include <string.h>
  51. #include <stdarg.h>
  52. static char *procname;
  53. static bool pinthread_override = false;
  54. static unsigned int ncore;
  55. static unsigned int setcore;
  56. static bool pinthread_debug = false;
  57. static cpu_set_t mask;
  58. static int (*real_pthread_create)(pthread_t *thread,
  59. const pthread_attr_t *attr,
  60. void *(*start_routine) (void *), void *arg);
  61. static int (*real_pthread_setaffinity_np)(pthread_t thread, size_t cpusetsize,
  62. const cpu_set_t *cpuset);
  63. static int (*real_pthread_attr_setaffinity_np)(pthread_attr_t *attr,
  64. size_t cpusetsize, const cpu_set_t *cpuset);
  65. static int (*real_sched_setaffinity)(pid_t pid, size_t cpusetsize,
  66. const cpu_set_t *mask);
  67. static void pint_debug(const char *msg, ...)
  68. {
  69. va_list args;
  70. char *str;
  71. int len;
  72. if(pinthread_debug)
  73. {
  74. va_start(args, msg);
  75. len = vsnprintf(NULL, 0, msg, args);
  76. if ((str = malloc(len+1)) == NULL)
  77. return;
  78. vsnprintf(str, len+1, msg, args);
  79. va_end(args);
  80. fprintf(stderr, "D:PINTHREAD: %s\n", str);
  81. free(str);
  82. }
  83. }
  84. static void main_pinthread(int argc, char* argv[], char* envp[])
  85. {
  86. char *pch;
  87. char *msg;
  88. CPU_ZERO(&mask); /* Initialize it all to 0, i.e. no CPUs selected. */
  89. procname = basename(argv[0]);
  90. msg = getenv("PINTHREAD_DEBUG");
  91. if(msg != NULL)
  92. pinthread_debug = true;
  93. msg = getenv("PINTHREAD_PNAMES");
  94. if(msg == NULL)
  95. {
  96. pinthread_override = true;
  97. } else {
  98. pch = strtok (msg," ");
  99. while (pch != NULL)
  100. {
  101. if (!strcmp(procname,pch))
  102. {
  103. pinthread_override = true;
  104. pch = NULL;
  105. } else {
  106. pch = strtok (NULL," ");
  107. }
  108. }
  109. }
  110. pint_debug("pinthread_main loaded");
  111. ncore = sysconf (_SC_NPROCESSORS_CONF);
  112. msg = getenv("PINTHREAD_CORE");
  113. if (msg != NULL)
  114. {
  115. setcore = (unsigned int) strtoul(msg, (char **)NULL, 10);
  116. pint_debug("Setting core as configured: %u", setcore);
  117. if(setcore >= ncore)
  118. {
  119. fprintf(stderr, "E:PINTHREAD wrong value for PINTHREAD_CORE: %u - using default.\n", setcore);
  120. setcore = sched_getcpu();
  121. }
  122. } else {
  123. setcore = sched_getcpu();
  124. pint_debug("Setting core by sched_getcpu: %u", setcore);
  125. }
  126. CPU_SET(setcore, &mask);
  127. pint_debug("Set real_sched_setaffinity");
  128. real_sched_setaffinity = dlsym(RTLD_NEXT,"sched_setaffinity");
  129. if ((msg=dlerror())!=NULL)
  130. fprintf(stderr, "E:PINTHREAD sched_setaffinity dlsym failed : %s\n", msg);
  131. // make sure the main thread is running on the same core:
  132. real_sched_setaffinity(getpid(), sizeof(mask), &mask);
  133. pint_debug("Affinity of main process configured");
  134. pint_debug("Set real_pthread_create");
  135. real_pthread_create = dlsym(RTLD_NEXT,"pthread_create");
  136. if ((msg=dlerror())!=NULL)
  137. fprintf(stderr, "E:PINTHREAD pthread_create dlsym failed : %s\n", msg);
  138. pint_debug("Set pthread_setaffinity_np");
  139. real_pthread_setaffinity_np = dlsym(RTLD_NEXT,"pthread_setaffinity_np");
  140. if ((msg=dlerror())!=NULL)
  141. fprintf(stderr, "E:PINTHREAD pthread_setaffinity_np dlsym failed : %s\n", msg);
  142. pint_debug("Set real_pthread_attr_setaffinity_np");
  143. real_pthread_attr_setaffinity_np = dlsym(RTLD_NEXT,"pthread_attr_setaffinity_np");
  144. if ((msg=dlerror())!=NULL)
  145. fprintf(stderr, "E:PINTHREAD pthread_attr_setaffinity_np dlsym failed : %s\n", msg);
  146. }
  147. __attribute__((section(".init_array"))) void (* p_main_pinthread)(int,char*[],char*[]) = &main_pinthread;
  148. int pthread_create(pthread_t *thread,
  149. const pthread_attr_t *attr,
  150. void *(*start_routine) (void *), void *arg)
  151. {
  152. int ret;
  153. pint_debug("about to call original pthread_create");
  154. ret = real_pthread_create(thread,attr,start_routine,arg);
  155. if(pinthread_override)
  156. {
  157. pint_debug("Overriding pthread_create");
  158. real_pthread_setaffinity_np(*thread, sizeof(mask), &mask);
  159. }
  160. return ret;
  161. }
  162. /*
  163. * We need to override also pthread_setaffinity family calls to manage cases in which
  164. * the running program try to change it by itself
  165. */
  166. int pthread_setaffinity_np(pthread_t thread, size_t cpusetsize,
  167. const cpu_set_t *cpuset)
  168. {
  169. pint_debug("about to call original pthread_setaffinity_np");
  170. if(pinthread_override)
  171. {
  172. pint_debug("Overriding pthread_setaffinity_np");
  173. return real_pthread_setaffinity_np(thread, sizeof(mask), &mask);
  174. }
  175. return real_pthread_setaffinity_np(thread, cpusetsize, cpuset);
  176. }
  177. int pthread_attr_setaffinity_np(pthread_attr_t *attr,
  178. size_t cpusetsize, const cpu_set_t *cpuset)
  179. {
  180. pint_debug("about to call original pthread_attr_setaffinity_np");
  181. if(pinthread_override)
  182. {
  183. pint_debug("Overriding pthread_attr_setaffinity_np");
  184. return real_pthread_attr_setaffinity_np(attr, sizeof(mask), &mask);
  185. }
  186. return real_pthread_attr_setaffinity_np(attr, cpusetsize, cpuset);
  187. }
  188. int sched_setaffinity(pid_t pid, size_t cpusetsize,
  189. const cpu_set_t *omask)
  190. {
  191. pint_debug("about to call original sched_setaffinity");
  192. if(pinthread_override)
  193. {
  194. pint_debug("Overriding sched_setaffinity");
  195. return real_sched_setaffinity(pid, sizeof(mask), &mask);
  196. }
  197. return real_sched_setaffinity(pid, cpusetsize, omask);
  198. }