aboutsummaryrefslogtreecommitdiff
path: root/toys/taskset.c
blob: c877264a0deb3f793e6cb232e65e5856ff5df0cb (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
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);
	}
}