diff options
| -rw-r--r-- | toys.h | 1 | ||||
| -rw-r--r-- | toys/taskset.c | 134 |
2 files changed, 135 insertions, 0 deletions
@@ -21,6 +21,7 @@ #include <pwd.h> #include <sched.h> #include <setjmp.h> +#include <sched.h> #include <shadow.h> #include <stdarg.h> #include <stdint.h> diff --git a/toys/taskset.c b/toys/taskset.c new file mode 100644 index 00000000..c877264a --- /dev/null +++ b/toys/taskset.c @@ -0,0 +1,134 @@ +/* vi: set sw=4 ts=4: + * + * taskset.c - Retrieve or set the CPU affinity of a process. + * + * Copyright 2012 Elie De Brauwer <eliedebrauwer@gmail.com> + * + * Not in SUSv4. + +USE_TASKSET(NEWTOY(taskset, "<1>2a", TOYFLAG_BIN|TOYFLAG_NEEDROOT)) + +config TASKSET + bool "taskset" + default y + help + usage: taskset [-a] [mask] PID + + When mask is present the CPU affinity mask of a given PID will + be set to this mask. When a mask is not given, the mask will + be printed. A mask is a hexadecimal string where the bit position + ..matches the cpu number. + -a Set/get the affinity of all tasks of a PID. +*/ + +#define _GNU_SOURCE +#include "toys.h" + +static int str_to_cpu_set(char * mask, cpu_set_t *set) +{ + int size = strlen(mask); + char *ptr = mask + size - 1; + int cpu = 0; + + CPU_ZERO(set); + if (size > 1 && mask[0] == '0' && mask[1] == 'x') + mask += 2; + + while(ptr >= mask) + { + char val = 0; + if ( *ptr >= '0' && *ptr <= '9') + val = *ptr - '0'; + else if (*ptr >= 'a' && *ptr <= 'f') + val = 10 + (*ptr - 'a'); + else + return -1; + + if (val & 1) CPU_SET(cpu, set); + if (val & 2) CPU_SET(cpu + 1, set); + if (val & 4) CPU_SET(cpu + 2, set); + if (val & 8) CPU_SET(cpu + 3, set); + + ptr--; + cpu += 4; + } + return 0; +} + +static char * cpu_set_to_str(cpu_set_t *set) +{ + int cpu; + char * ptr = toybuf; + for (cpu=(8*sizeof(cpu_set_t) - 4); cpu >= 0; cpu -= 4) + { + char val = 0; + if (CPU_ISSET(cpu, set)) val |= 1; + if (CPU_ISSET(cpu + 1, set)) val |= 2; + if (CPU_ISSET(cpu + 2, set)) val |= 4; + if (CPU_ISSET(cpu + 3, set)) val |= 8; + if (ptr != toybuf || val != 0) + { + if (val < 10) + *ptr = '0' + val; + else + *ptr = 'a' + (val - 10); + ptr++; + } + } + *ptr = 0; + return toybuf; +} + +static void do_taskset(pid_t pid) +{ + cpu_set_t mask; + if (sched_getaffinity(pid, sizeof(mask), &mask)) + perror_exit("failed to get %d's affinity", pid); + + printf("pid %d's current affinity mask: %s\n", pid, cpu_set_to_str(&mask)); + + if (toys.optc == 2) + { + if (str_to_cpu_set(toys.optargs[0], &mask)) + perror_exit("failed to parse CPU mask: %s", toys.optargs[0]); + + if (sched_setaffinity(pid, sizeof(mask), &mask)) + perror_exit("failed to set %d's affinity", pid); + + if (sched_getaffinity(pid, sizeof(mask), &mask)) + perror_exit("failed to get %d's affinity", pid); + + printf("pid %d's new affinity mask: %s\n", pid, cpu_set_to_str(&mask)); + } +} + +static int task_cb(struct dirtree *new) +{ + if (S_ISDIR(new->st.st_mode)) + { + if (strstr(new->name,"/task/\0") != NULL) + return DIRTREE_RECURSE; + + pid_t tid = atoi(new->name); + if (tid) + do_taskset(tid); + } + return 0; +} + +void taskset_main(void) +{ + char * pidstr = (toys.optc==1)?toys.optargs[0]:toys.optargs[1]; + + if ( toys.optflags & 0x1 ) + { + sprintf(toybuf, "/proc/%s/task/", pidstr); + dirtree_read(toybuf, task_cb); + } + else + { + pid_t tid = atoi(pidstr); + if (tid) + do_taskset(tid); + } +} |