1/*
2 * taskset.c - taskset
3 * Command-line utility for setting and retrieving a task's CPU affinity
4 *
5 * Robert Love <rml@tech9.net>		25 April 2002
6 *
7 * Linux kernels as of 2.5.8 provide the needed syscalls for
8 * working with a task's cpu affinity.  Currently 2.4 does not
9 * support these syscalls, but patches are available at:
10 *
11 * 	http://www.kernel.org/pub/linux/kernel/people/rml/cpu-affinity/
12 *
13 * This program is free software; you can redistribute it and/or modify
14 * it under the terms of the GNU General Public License, v2, as
15 * published by the Free Software Foundation
16 *
17 * This program is distributed in the hope that it will be useful,
18 * but WITHOUT ANY WARRANTY; without even the implied warranty of
19 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
20 * GNU General Public License for more details.
21 *
22 * You should have received a copy of the GNU General Public License
23 * along with this program; if not, write to the Free Software
24 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
25 *
26 * Copyright (C) 2004 Robert Love
27 */
28
29#define _GNU_SOURCE
30
31#include <stdio.h>
32#include <stdlib.h>
33#include <unistd.h>
34#include <getopt.h>
35#include <sched.h>
36#include <errno.h>
37#include <string.h>
38#include <ctype.h>
39
40static void show_usage(const char *cmd)
41{
42	fprintf(stderr, "taskset version " "VERSION" "\n");
43	fprintf(stderr, "usage: %s [options] [mask | cpu-list] [pid |"\
44		" cmd [args...]]\n", cmd);
45	fprintf(stderr, "set or get the affinity of a process\n\n");
46	fprintf(stderr, "  -p, --pid                  "
47		"operate on existing given pid\n");
48        fprintf(stderr, "  -c, --cpu-list             "\
49		"display and specify cpus in list format\n");
50	fprintf(stderr, "  -h, --help                 display this help\n");
51	fprintf(stderr, "  -v, --version              "\
52		"output version information\n\n");
53	fprintf(stderr, "The default behavior is to run a new command:\n");
54	fprintf(stderr, "  %s 03 sshd -b 1024\n", cmd);
55	fprintf(stderr, "You can retrieve the mask of an existing task:\n");
56	fprintf(stderr, "  %s -p 700\n", cmd);
57	fprintf(stderr, "Or set it:\n");
58	fprintf(stderr, "  %s -p 03 700\n", cmd);
59	fprintf(stderr, "List format uses a comma-separated list instead"\
60			" of a mask:\n");
61	fprintf(stderr, "  %s -pc 0,3,7-11 700\n", cmd);
62	fprintf(stderr, "Ranges in list format can take a stride argument:\n");
63	fprintf(stderr, "  e.g. 0-31:2 is equivalent to mask 0x55555555\n\n");
64}
65
66static inline int val_to_char(int v)
67{
68	if (v >= 0 && v < 10)
69		return '0' + v;
70	else if (v >= 10 && v < 16)
71		return ('a' - 10) + v;
72	else
73		return -1;
74}
75
76static char * cpuset_to_str(cpu_set_t *mask, char *str)
77{
78	int base;
79	char *ptr = str;
80	char *ret = 0;
81
82	for (base = CPU_SETSIZE - 4; base >= 0; base -= 4) {
83		char val = 0;
84		if (CPU_ISSET(base, mask))
85			val |= 1;
86		if (CPU_ISSET(base + 1, mask))
87			val |= 2;
88		if (CPU_ISSET(base + 2, mask))
89			val |= 4;
90		if (CPU_ISSET(base + 3, mask))
91			val |= 8;
92		if (!ret && val)
93			ret = ptr;
94		*ptr++ = val_to_char(val);
95	}
96	*ptr = 0;
97	return ret ? ret : ptr - 1;
98}
99
100static char * cpuset_to_cstr(cpu_set_t *mask, char *str)
101{
102	int i;
103	char *ptr = str;
104	int entry_made = 0;
105
106	for (i = 0; i < CPU_SETSIZE; i++) {
107		if (CPU_ISSET(i, mask)) {
108			int j;
109			int run = 0;
110			entry_made = 1;
111			for (j = i + 1; j < CPU_SETSIZE; j++) {
112				if (CPU_ISSET(j, mask))
113					run++;
114				else
115					break;
116			}
117			if (!run)
118				sprintf(ptr, "%d,", i);
119			else if (run == 1) {
120				sprintf(ptr, "%d,%d,", i, i + 1);
121				i++;
122			} else {
123				sprintf(ptr, "%d-%d,", i, i + run);
124				i += run;
125			}
126			while (*ptr != 0)
127				ptr++;
128		}
129	}
130	ptr -= entry_made;
131	*ptr = 0;
132
133	return str;
134}
135
136static inline int char_to_val(int c)
137{
138	int cl;
139
140	cl = tolower(c);
141	if (c >= '0' && c <= '9')
142		return c - '0';
143	else if (cl >= 'a' && cl <= 'f')
144		return cl + (10 - 'a');
145	else
146		return -1;
147}
148
149static int str_to_cpuset(cpu_set_t *mask, const char* str)
150{
151	int len = strlen(str);
152	const char *ptr = str + len - 1;
153	int base = 0;
154
155	/* skip 0x, it's all hex anyway */
156	if (len > 1 && !memcmp(str, "0x", 2L))
157		str += 2;
158
159	CPU_ZERO(mask);
160	while (ptr >= str) {
161		char val = char_to_val(*ptr);
162		if (val == (char) -1)
163			return -1;
164		if (val & 1)
165			CPU_SET(base, mask);
166		if (val & 2)
167			CPU_SET(base + 1, mask);
168		if (val & 4)
169			CPU_SET(base + 2, mask);
170		if (val & 8)
171			CPU_SET(base + 3, mask);
172		len--;
173		ptr--;
174		base += 4;
175	}
176
177	return 0;
178}
179
180static const char *nexttoken(const char *q,  int sep)
181{
182	if (q)
183		q = strchr(q, sep);
184	if (q)
185		q++;
186	return q;
187}
188
189static int cstr_to_cpuset(cpu_set_t *mask, const char* str)
190{
191	const char *p, *q;
192	q = str;
193	CPU_ZERO(mask);
194
195	while (p = q, q = nexttoken(q, ','), p) {
196		unsigned int a;	/* beginning of range */
197		unsigned int b;	/* end of range */
198		unsigned int s;	/* stride */
199		const char *c1, *c2;
200
201		if (sscanf(p, "%u", &a) < 1)
202			return 1;
203		b = a;
204		s = 1;
205
206		c1 = nexttoken(p, '-');
207		c2 = nexttoken(p, ',');
208		if (c1 != NULL && (c2 == NULL || c1 < c2)) {
209			if (sscanf(c1, "%u", &b) < 1)
210				return 1;
211			c1 = nexttoken(c1, ':');
212			if (c1 != NULL && (c2 == NULL || c1 < c2))
213				if (sscanf(c1, "%u", &s) < 1) {
214					return 1;
215			}
216		}
217
218		if (!(a <= b))
219			return 1;
220		while (a <= b) {
221			CPU_SET(a, mask);
222			a += s;
223		}
224	}
225
226	return 0;
227}
228
229int main(int argc, char *argv[])
230{
231	cpu_set_t new_mask, cur_mask;
232	pid_t pid = 0;
233	int opt, err;
234	char mstr[1 + CPU_SETSIZE / 4];
235	char cstr[7 * CPU_SETSIZE];
236	int c_opt = 0;
237
238	struct option longopts[] = {
239		{ "pid",	0, NULL, 'p' },
240		{ "cpu-list",	0, NULL, 'c' },
241		{ "help",	0, NULL, 'h' },
242		{ "version",	0, NULL, 'V' },
243		{ NULL,		0, NULL, 0 }
244	};
245
246	while ((opt = getopt_long(argc, argv, "+pchV", longopts, NULL)) != -1) {
247		int ret = 1;
248
249		switch (opt) {
250		case 'p':
251			pid = atoi(argv[argc - 1]);
252			break;
253		case 'c':
254			c_opt = 1;
255			break;
256		case 'V':
257			printf("taskset version " "VERSION" "\n");
258			return 0;
259		case 'h':
260			ret = 0;
261		default:
262			show_usage(argv[0]);
263			return ret;
264		}
265	}
266
267	if ((!pid && argc - optind < 2)
268			|| (pid && (argc - optind < 1 || argc - optind > 2))) {
269		show_usage(argv[0]);
270		return 1;
271	}
272
273	if (pid) {
274		if (sched_getaffinity(pid, sizeof (cur_mask), &cur_mask) < 0) {
275			perror("sched_getaffinity");
276			fprintf(stderr, "failed to get pid %d's affinity\n",
277				pid);
278			return 1;
279		}
280		if (c_opt)
281			printf("pid %d's current affinity list: %s\n", pid,
282				cpuset_to_cstr(&cur_mask, cstr));
283		else
284			printf("pid %d's current affinity mask: %s\n", pid,
285				cpuset_to_str(&cur_mask, mstr));
286
287		if (argc - optind == 1)
288			return 0;
289	}
290
291	if (c_opt)
292		err = cstr_to_cpuset(&new_mask, argv[optind]);
293	else
294		err = str_to_cpuset(&new_mask, argv[optind]);
295
296	if (err) {
297		if (c_opt)
298			fprintf(stderr, "failed to parse CPU list %s\n",
299				argv[optind]);
300		else
301			fprintf(stderr, "failed to parse CPU mask %s\n",
302				argv[optind]);
303		return 1;
304	}
305
306	if (sched_setaffinity(pid, sizeof (new_mask), &new_mask)) {
307		perror("sched_setaffinity");
308		fprintf(stderr, "failed to set pid %d's affinity.\n", pid);
309		return 1;
310	}
311
312	if (sched_getaffinity(pid, sizeof (cur_mask), &cur_mask) < 0) {
313		perror("sched_getaffinity");
314		fprintf(stderr, "failed to get pid %d's affinity.\n", pid);
315		return 1;
316	}
317
318	if (pid) {
319		if (c_opt)
320			printf("pid %d's new affinity list: %s\n", pid,
321		               cpuset_to_cstr(&cur_mask, cstr));
322		else
323			printf("pid %d's new affinity mask: %s\n", pid,
324		               cpuset_to_str(&cur_mask, mstr));
325	} else {
326		argv += optind + 1;
327		execvp(argv[0], argv);
328		perror("execvp");
329		fprintf(stderr, "failed to execute %s\n", argv[0]);
330		return 1;
331	}
332
333	return 0;
334}
335