cpuset.c revision 336039
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: stable/10/usr.bin/cpuset/cpuset.c 336039 2018-07-06 19:10:07Z jamie $");
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 <jail.h>
43#include <limits.h>
44#include <stdio.h>
45#include <stdlib.h>
46#include <stdint.h>
47#include <unistd.h>
48#include <string.h>
49
50static int Cflag;
51static int cflag;
52static int gflag;
53static int iflag;
54static int jflag;
55static int lflag;
56static int pflag;
57static int rflag;
58static int sflag;
59static int tflag;
60static int xflag;
61static id_t id;
62static cpulevel_t level;
63static cpuwhich_t which;
64
65static void usage(void);
66
67static void printset(cpuset_t *mask);
68
69static void
70parselist(char *list, cpuset_t *mask)
71{
72	enum { NONE, NUM, DASH } state;
73	int lastnum;
74	int curnum;
75	char *l;
76
77	if (strcasecmp(list, "all") == 0) {
78		if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
79		    sizeof(*mask), mask) != 0)
80			err(EXIT_FAILURE, "getaffinity");
81		return;
82	}
83	state = NONE;
84	curnum = lastnum = 0;
85	for (l = list; *l != '\0';) {
86		if (isdigit(*l)) {
87			curnum = atoi(l);
88			if (curnum > CPU_SETSIZE)
89				errx(EXIT_FAILURE,
90				    "Only %d cpus supported", CPU_SETSIZE);
91			while (isdigit(*l))
92				l++;
93			switch (state) {
94			case NONE:
95				lastnum = curnum;
96				state = NUM;
97				break;
98			case DASH:
99				for (; lastnum <= curnum; lastnum++)
100					CPU_SET(lastnum, mask);
101				state = NONE;
102				break;
103			case NUM:
104			default:
105				goto parserr;
106			}
107			continue;
108		}
109		switch (*l) {
110		case ',':
111			switch (state) {
112			case NONE:
113				break;
114			case NUM:
115				CPU_SET(curnum, mask);
116				state = NONE;
117				break;
118			case DASH:
119				goto parserr;
120				break;
121			}
122			break;
123		case '-':
124			if (state != NUM)
125				goto parserr;
126			state = DASH;
127			break;
128		default:
129			goto parserr;
130		}
131		l++;
132	}
133	switch (state) {
134		case NONE:
135			break;
136		case NUM:
137			CPU_SET(curnum, mask);
138			break;
139		case DASH:
140			goto parserr;
141	}
142	return;
143parserr:
144	errx(EXIT_FAILURE, "Malformed cpu-list %s", list);
145}
146
147static void
148printset(cpuset_t *mask)
149{
150	int once;
151	int cpu;
152
153	for (once = 0, cpu = 0; cpu < CPU_SETSIZE; cpu++) {
154		if (CPU_ISSET(cpu, mask)) {
155			if (once == 0) {
156				printf("%d", cpu);
157				once = 1;
158			} else
159				printf(", %d", cpu);
160		}
161	}
162	printf("\n");
163}
164
165static const char *whichnames[] = { NULL, "tid", "pid", "cpuset", "irq", "jail" };
166static const char *levelnames[] = { NULL, " root", " cpuset", "" };
167
168static void
169printaffinity(void)
170{
171	cpuset_t mask;
172
173	if (cpuset_getaffinity(level, which, id, sizeof(mask), &mask) != 0)
174		err(EXIT_FAILURE, "getaffinity");
175	printf("%s %jd%s mask: ", whichnames[which], (intmax_t)id,
176	    levelnames[level]);
177	printset(&mask);
178	exit(EXIT_SUCCESS);
179}
180
181static void
182printsetid(void)
183{
184	cpusetid_t setid;
185
186	/*
187	 * Only LEVEL_WHICH && WHICH_CPUSET has a numbered id.
188	 */
189	if (level == CPU_LEVEL_WHICH && !sflag)
190		level = CPU_LEVEL_CPUSET;
191	if (cpuset_getid(level, which, id, &setid))
192		err(errno, "getid");
193	printf("%s %jd%s id: %d\n", whichnames[which], (intmax_t)id,
194	    levelnames[level], setid);
195}
196
197int
198main(int argc, char *argv[])
199{
200	cpusetid_t setid;
201	cpuset_t mask;
202	lwpid_t tid;
203	pid_t pid;
204	int ch;
205
206	CPU_ZERO(&mask);
207	level = CPU_LEVEL_WHICH;
208	which = CPU_WHICH_PID;
209	id = pid = tid = setid = -1;
210	while ((ch = getopt(argc, argv, "Ccgij:l:p:rs:t:x:")) != -1) {
211		switch (ch) {
212		case 'C':
213			Cflag = 1;
214			break;
215		case 'c':
216			if (rflag)
217				usage();
218			cflag = 1;
219			level = CPU_LEVEL_CPUSET;
220			break;
221		case 'g':
222			gflag = 1;
223			break;
224		case 'i':
225			iflag = 1;
226			break;
227		case 'j':
228			jflag = 1;
229			which = CPU_WHICH_JAIL;
230			id = jail_getid(optarg);
231			if (id < 0)
232				errx(EXIT_FAILURE, "%s", jail_errmsg);
233			break;
234		case 'l':
235			lflag = 1;
236			parselist(optarg, &mask);
237			break;
238		case 'p':
239			pflag = 1;
240			which = CPU_WHICH_PID;
241			id = pid = atoi(optarg);
242			break;
243		case 'r':
244			if (cflag)
245				usage();
246			level = CPU_LEVEL_ROOT;
247			rflag = 1;
248			break;
249		case 's':
250			sflag = 1;
251			which = CPU_WHICH_CPUSET;
252			id = setid = atoi(optarg);
253			break;
254		case 't':
255			tflag = 1;
256			which = CPU_WHICH_TID;
257			id = tid = atoi(optarg);
258			break;
259		case 'x':
260			xflag = 1;
261			which = CPU_WHICH_IRQ;
262			id = atoi(optarg);
263			break;
264		default:
265			usage();
266		}
267	}
268	argc -= optind;
269	argv += optind;
270	if (gflag) {
271		if (argc || Cflag || lflag)
272			usage();
273		/* Only one identity specifier. */
274		if (jflag + xflag + sflag + pflag + tflag > 1)
275			usage();
276		if (iflag)
277			printsetid();
278		else
279			printaffinity();
280		exit(EXIT_SUCCESS);
281	}
282	if (iflag)
283		usage();
284	/*
285	 * The user wants to run a command with a set and possibly cpumask.
286	 */
287	if (argc) {
288		if (Cflag | pflag | rflag | tflag | xflag | jflag)
289			usage();
290		if (sflag) {
291			if (cpuset_setid(CPU_WHICH_PID, -1, setid))
292				err(argc, "setid");
293		} else {
294			if (cpuset(&setid))
295				err(argc, "newid");
296		}
297		if (lflag) {
298			if (cpuset_setaffinity(level, CPU_WHICH_PID,
299			    -1, sizeof(mask), &mask) != 0)
300				err(EXIT_FAILURE, "setaffinity");
301		}
302		errno = 0;
303		execvp(*argv, argv);
304		err(errno == ENOENT ? 127 : 126, "%s", *argv);
305	}
306	/*
307	 * We're modifying something that presently exists.
308	 */
309	if (Cflag && (sflag || rflag || !pflag || tflag || xflag || jflag))
310		usage();
311	if (!lflag && (cflag || rflag))
312		usage();
313	if (!lflag && !(Cflag || sflag))
314		usage();
315	/* You can only set a mask on a thread. */
316	if (tflag && (sflag | pflag | xflag | jflag))
317		usage();
318	/* You can only set a mask on an irq. */
319	if (xflag && (jflag | pflag | sflag | tflag))
320		usage();
321	if (Cflag) {
322		/*
323		 * Create a new cpuset and move the specified process
324		 * into the set.
325		 */
326		if (cpuset(&setid) < 0)
327			err(EXIT_FAILURE, "newid");
328		sflag = 1;
329	}
330	if (pflag && sflag) {
331		if (cpuset_setid(CPU_WHICH_PID, pid, setid))
332			err(EXIT_FAILURE, "setid");
333		/*
334		 * If the user specifies a set and a list we want the mask
335		 * to effect the pid and not the set.
336		 */
337		which = CPU_WHICH_PID;
338		id = pid;
339	}
340	if (lflag) {
341		if (cpuset_setaffinity(level, which, id, sizeof(mask),
342		    &mask) != 0)
343			err(EXIT_FAILURE, "setaffinity");
344	}
345
346	exit(EXIT_SUCCESS);
347}
348
349static void
350usage(void)
351{
352
353	fprintf(stderr,
354	    "usage: cpuset [-l cpu-list] [-s setid] cmd ...\n");
355	fprintf(stderr,
356	    "       cpuset [-l cpu-list] [-s setid] -p pid\n");
357	fprintf(stderr,
358	    "       cpuset [-c] [-l cpu-list] -C -p pid\n");
359	fprintf(stderr,
360	    "       cpuset [-cr] [-l cpu-list] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
361	fprintf(stderr,
362	    "       cpuset [-cgir] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
363	exit(1);
364}
365