1 /* taskset.c - Retrieve or set the CPU affinity of a process. 2 * 3 * Copyright 2012 Elie De Brauwer <eliedebrauwer (at) gmail.com> 4 5 USE_TASKSET(NEWTOY(taskset, "<1^pa", TOYFLAG_BIN|TOYFLAG_STAYROOT)) 6 USE_NPROC(NEWTOY(nproc, "(all)", TOYFLAG_USR|TOYFLAG_BIN)) 7 8 config NPROC 9 bool "nproc" 10 default y 11 help 12 usage: nproc [--all] 13 14 Print number of processors. 15 16 --all Show all processors, not just ones this task can run on 17 18 config TASKSET 19 bool "taskset" 20 default y 21 help 22 usage: taskset [-ap] [mask] [PID | cmd [args...]] 23 24 Launch a new task which may only run on certain processors, or change 25 the processor affinity of an exisitng PID. 26 27 Mask is a hex string where each bit represents a processor the process 28 is allowed to run on. PID without a mask displays existing affinity. 29 30 -p Set/get the affinity of given PID instead of a new command 31 -a Set/get the affinity of all threads of the PID 32 */ 33 34 #define FOR_taskset 35 #include "toys.h" 36 37 #include <sys/syscall.h> 38 #define sched_setaffinity(pid, size, cpuset) \ 39 syscall(__NR_sched_setaffinity, (pid_t)pid, (size_t)size, (void *)cpuset) 40 #define sched_getaffinity(pid, size, cpuset) \ 41 syscall(__NR_sched_getaffinity, (pid_t)pid, (size_t)size, (void *)cpuset) 42 43 // mask is an array of long, which makes the layout a bit weird on big 44 // endian systems but as long as it's consistent... 45 46 static void do_taskset(pid_t pid, int quiet) 47 { 48 unsigned long *mask = (unsigned long *)toybuf; 49 char *s = *toys.optargs, *failed = "failed to %s %d's affinity"; 50 int i, j, k; 51 52 for (i=0; ; i++) { 53 if (!quiet) { 54 int j = sizeof(toybuf), flag = 0; 55 56 if (-1 == sched_getaffinity(pid, sizeof(toybuf), (void *)mask)) 57 perror_exit(failed, "get", pid); 58 59 printf("pid %d's %s affinity mask: ", pid, i ? "new" : "current"); 60 61 while (j--) { 62 int x = 255 & (mask[j/sizeof(long)] >> (8*(j&(sizeof(long)-1)))); 63 64 if (flag) printf("%02x", x); 65 else if (x) { 66 flag++; 67 printf("%x", x); 68 } 69 } 70 putchar('\n'); 71 } 72 73 if (i || toys.optc < 2) return; 74 75 memset(toybuf, 0, sizeof(toybuf)); 76 k = strlen(s = *toys.optargs); 77 s += k; 78 for (j = 0; j<k; j++) { 79 unsigned long digit = *(--s) - '0'; 80 81 if (digit > 9) digit = 10 + tolower(*s)-'a'; 82 if (digit > 15) error_exit("bad mask '%s'", *toys.optargs); 83 mask[j/(2*sizeof(long))] |= digit << 4*(j&((2*sizeof(long))-1)); 84 } 85 86 if (-1 == sched_setaffinity(pid, sizeof(toybuf), (void *)mask)) 87 perror_exit(failed, "set", pid); 88 } 89 } 90 91 static int task_callback(struct dirtree *new) 92 { 93 if (!new->parent) return DIRTREE_RECURSE; 94 if (isdigit(*new->name)) do_taskset(atoi(new->name), 0); 95 96 return 0; 97 } 98 99 void taskset_main(void) 100 { 101 if (!(toys.optflags & FLAG_p)) { 102 if (toys.optc < 2) error_exit("Needs 2 args"); 103 do_taskset(getpid(), 1); 104 xexec(toys.optargs+1); 105 } else { 106 char *c; 107 pid_t pid = strtol(toys.optargs[toys.optc-1], &c, 10); 108 109 if (*c) error_exit("Not int %s", toys.optargs[1]); 110 111 if (toys.optflags & FLAG_a) { 112 char buf[33]; 113 sprintf(buf, "/proc/%ld/task/", (long)pid); 114 dirtree_read(buf, task_callback); 115 } else do_taskset(pid, 0); 116 } 117 } 118 119 void nproc_main(void) 120 { 121 unsigned i, j, nproc = 0; 122 123 // This can only detect 32768 processors. Call getaffinity and count bits. 124 if (!toys.optflags && -1!=sched_getaffinity(getpid(), 4096, toybuf)) { 125 for (i = 0; i<4096; i++) 126 if (toybuf[i]) for (j=0; j<8; j++) if (toybuf[i]&(1<<j)) nproc++; 127 } 128 129 // If getaffinity failed or --all, count cpu entries in proc 130 if (!nproc) nproc = sysconf(_SC_NPROCESSORS_CONF); 131 132 xprintf("%u\n", nproc); 133 } 134
