1/*
2 * Copyright (c) 2007, 2008 	Jeffrey Roberson <jeff@freebsd.org>
3 * All rights reserved.
4 *
5 * Copyright (c) 2008 Nokia Corporation
6 * All rights reserved.
7 *
8 * Redistribution and use in source and binary forms, with or without
9 * modification, are permitted provided that the following conditions
10 * are met:
11 * 1. Redistributions of source code must retain the above copyright
12 *    notice, this list of conditions and the following disclaimer.
13 * 2. Redistributions in binary form must reproduce the above copyright
14 *    notice, this list of conditions and the following disclaimer in the
15 *    documentation and/or other materials provided with the distribution.
16 *
17 * THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
18 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
19 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
20 * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
21 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
22 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
23 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
24 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
25 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
26 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
27 * SUCH DAMAGE.
28 */
29
30#include <sys/cdefs.h>
31__FBSDID("$FreeBSD$");
32
33#include <sys/param.h>
34#include <sys/types.h>
35#include <sys/time.h>
36#include <sys/resource.h>
37#include <sys/cpuset.h>
38
39#include <ctype.h>
40#include <err.h>
41#include <errno.h>
42#include <limits.h>
43#include <stdio.h>
44#include <stdlib.h>
45#include <stdint.h>
46#include <unistd.h>
47#include <string.h>
48
49static int Cflag;
50static int cflag;
51static int gflag;
52static int iflag;
53static int jflag;
54static int lflag;
55static int pflag;
56static int rflag;
57static int sflag;
58static int tflag;
59static int xflag;
60static id_t id;
61static cpulevel_t level;
62static cpuwhich_t which;
63
64static void usage(void);
65
66static void printset(cpuset_t *mask);
67
68static void
69parselist(char *list, cpuset_t *mask)
70{
71	enum { NONE, NUM, DASH } state;
72	int lastnum;
73	int curnum;
74	char *l;
75
76	if (strcasecmp(list, "all") == 0) {
77		if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
78		    sizeof(*mask), mask) != 0)
79			err(EXIT_FAILURE, "getaffinity");
80		return;
81	}
82	state = NONE;
83	curnum = lastnum = 0;
84	for (l = list; *l != '\0';) {
85		if (isdigit(*l)) {
86			curnum = atoi(l);
87			if (curnum > CPU_SETSIZE)
88				errx(EXIT_FAILURE,
89				    "Only %d cpus supported", CPU_SETSIZE);
90			while (isdigit(*l))
91				l++;
92			switch (state) {
93			case NONE:
94				lastnum = curnum;
95				state = NUM;
96				break;
97			case DASH:
98				for (; lastnum <= curnum; lastnum++)
99					CPU_SET(lastnum, mask);
100				state = NONE;
101				break;
102			case NUM:
103			default:
104				goto parserr;
105			}
106			continue;
107		}
108		switch (*l) {
109		case ',':
110			switch (state) {
111			case NONE:
112				break;
113			case NUM:
114				CPU_SET(curnum, mask);
115				state = NONE;
116				break;
117			case DASH:
118				goto parserr;
119				break;
120			}
121			break;
122		case '-':
123			if (state != NUM)
124				goto parserr;
125			state = DASH;
126			break;
127		default:
128			goto parserr;
129		}
130		l++;
131	}
132	switch (state) {
133		case NONE:
134			break;
135		case NUM:
136			CPU_SET(curnum, mask);
137			break;
138		case DASH:
139			goto parserr;
140	}
141	return;
142parserr:
143	errx(EXIT_FAILURE, "Malformed cpu-list %s", list);
144}
145
146static void
147printset(cpuset_t *mask)
148{
149	int once;
150	int cpu;
151
152	for (once = 0, cpu = 0; cpu < CPU_SETSIZE; cpu++) {
153		if (CPU_ISSET(cpu, mask)) {
154			if (once == 0) {
155				printf("%d", cpu);
156				once = 1;
157			} else
158				printf(", %d", cpu);
159		}
160	}
161	printf("\n");
162}
163
164static const char *whichnames[] = { NULL, "tid", "pid", "cpuset", "irq", "jail" };
165static const char *levelnames[] = { NULL, " root", " cpuset", "" };
166
167static void
168printaffinity(void)
169{
170	cpuset_t mask;
171
172	if (cpuset_getaffinity(level, which, id, sizeof(mask), &mask) != 0)
173		err(EXIT_FAILURE, "getaffinity");
174	printf("%s %jd%s mask: ", whichnames[which], (intmax_t)id,
175	    levelnames[level]);
176	printset(&mask);
177	exit(EXIT_SUCCESS);
178}
179
180static void
181printsetid(void)
182{
183	cpusetid_t setid;
184
185	/*
186	 * Only LEVEL_WHICH && WHICH_CPUSET has a numbered id.
187	 */
188	if (level == CPU_LEVEL_WHICH && !sflag)
189		level = CPU_LEVEL_CPUSET;
190	if (cpuset_getid(level, which, id, &setid))
191		err(errno, "getid");
192	printf("%s %jd%s id: %d\n", whichnames[which], (intmax_t)id,
193	    levelnames[level], setid);
194}
195
196int
197main(int argc, char *argv[])
198{
199	cpusetid_t setid;
200	cpuset_t mask;
201	lwpid_t tid;
202	pid_t pid;
203	int ch;
204
205	CPU_ZERO(&mask);
206	level = CPU_LEVEL_WHICH;
207	which = CPU_WHICH_PID;
208	id = pid = tid = setid = -1;
209	while ((ch = getopt(argc, argv, "Ccgij:l:p:rs:t:x:")) != -1) {
210		switch (ch) {
211		case 'C':
212			Cflag = 1;
213			break;
214		case 'c':
215			if (rflag)
216				usage();
217			cflag = 1;
218			level = CPU_LEVEL_CPUSET;
219			break;
220		case 'g':
221			gflag = 1;
222			break;
223		case 'i':
224			iflag = 1;
225			break;
226		case 'j':
227			jflag = 1;
228			which = CPU_WHICH_JAIL;
229			id = atoi(optarg);
230			break;
231		case 'l':
232			lflag = 1;
233			parselist(optarg, &mask);
234			break;
235		case 'p':
236			pflag = 1;
237			which = CPU_WHICH_PID;
238			id = pid = atoi(optarg);
239			break;
240		case 'r':
241			if (cflag)
242				usage();
243			level = CPU_LEVEL_ROOT;
244			rflag = 1;
245			break;
246		case 's':
247			sflag = 1;
248			which = CPU_WHICH_CPUSET;
249			id = setid = atoi(optarg);
250			break;
251		case 't':
252			tflag = 1;
253			which = CPU_WHICH_TID;
254			id = tid = atoi(optarg);
255			break;
256		case 'x':
257			xflag = 1;
258			which = CPU_WHICH_IRQ;
259			id = atoi(optarg);
260			break;
261		default:
262			usage();
263		}
264	}
265	argc -= optind;
266	argv += optind;
267	if (gflag) {
268		if (argc || Cflag || lflag)
269			usage();
270		/* Only one identity specifier. */
271		if (jflag + xflag + sflag + pflag + tflag > 1)
272			usage();
273		if (iflag)
274			printsetid();
275		else
276			printaffinity();
277		exit(EXIT_SUCCESS);
278	}
279	if (iflag)
280		usage();
281	/*
282	 * The user wants to run a command with a set and possibly cpumask.
283	 */
284	if (argc) {
285		if (Cflag | pflag | rflag | tflag | xflag | jflag)
286			usage();
287		if (sflag) {
288			if (cpuset_setid(CPU_WHICH_PID, -1, setid))
289				err(argc, "setid");
290		} else {
291			if (cpuset(&setid))
292				err(argc, "newid");
293		}
294		if (lflag) {
295			if (cpuset_setaffinity(level, CPU_WHICH_PID,
296			    -1, sizeof(mask), &mask) != 0)
297				err(EXIT_FAILURE, "setaffinity");
298		}
299		errno = 0;
300		execvp(*argv, argv);
301		err(errno == ENOENT ? 127 : 126, "%s", *argv);
302	}
303	/*
304	 * We're modifying something that presently exists.
305	 */
306	if (Cflag && (sflag || rflag || !pflag || tflag || xflag || jflag))
307		usage();
308	if (!lflag && (cflag || rflag))
309		usage();
310	if (!lflag && !(Cflag || sflag))
311		usage();
312	/* You can only set a mask on a thread. */
313	if (tflag && (sflag | pflag | xflag | jflag))
314		usage();
315	/* You can only set a mask on an irq. */
316	if (xflag && (jflag | pflag | sflag | tflag))
317		usage();
318	if (Cflag) {
319		/*
320		 * Create a new cpuset and move the specified process
321		 * into the set.
322		 */
323		if (cpuset(&setid) < 0)
324			err(EXIT_FAILURE, "newid");
325		sflag = 1;
326	}
327	if (pflag && sflag) {
328		if (cpuset_setid(CPU_WHICH_PID, pid, setid))
329			err(EXIT_FAILURE, "setid");
330		/*
331		 * If the user specifies a set and a list we want the mask
332		 * to effect the pid and not the set.
333		 */
334		which = CPU_WHICH_PID;
335		id = pid;
336	}
337	if (lflag) {
338		if (cpuset_setaffinity(level, which, id, sizeof(mask),
339		    &mask) != 0)
340			err(EXIT_FAILURE, "setaffinity");
341	}
342
343	exit(EXIT_SUCCESS);
344}
345
346static void
347usage(void)
348{
349
350	fprintf(stderr,
351	    "usage: cpuset [-l cpu-list] [-s setid] cmd ...\n");
352	fprintf(stderr,
353	    "       cpuset [-l cpu-list] [-s setid] -p pid\n");
354	fprintf(stderr,
355	    "       cpuset [-c] [-l cpu-list] -C -p pid\n");
356	fprintf(stderr,
357	    "       cpuset [-cr] [-l cpu-list] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
358	fprintf(stderr,
359	    "       cpuset [-cgir] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
360	exit(1);
361}
362