Adding initial implementation of taskset

diff --git a/toys.h b/toys.h
index a333986..a03ed7c 100644 (file)
--- a/toys.h
+++ b/toys.h
@@ -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 (file)
index 0000000..c877264
--- /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);
+       }
+}