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
49 static int Cflag;
50 static int cflag;
51 static int dflag;
52 static int gflag;
53 static int iflag;
54 static int jflag;
55 static int lflag;
56 static int pflag;
57 static int rflag;
58 static int sflag;
59 static int tflag;
60 static int xflag;
61 static id_t id;
62 static cpulevel_t level;
63 static cpuwhich_t which;
64
65 static void usage(void);
66
67 static void printset(cpuset_t *mask);
68
69 static void
parselist(char * list,cpuset_t * mask)70 parselist(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;
143 parserr:
144 errx(EXIT_FAILURE, "Malformed cpu-list %s", list);
145 }
146
147 static void
printset(cpuset_t * mask)148 printset(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
165 static const char *whichnames[] = { NULL, "tid", "pid", "cpuset", "irq", "jail",
166 "domain" };
167 static const char *levelnames[] = { NULL, " root", " cpuset", "" };
168
169 static void
printaffinity(void)170 printaffinity(void)
171 {
172 cpuset_t mask;
173
174 if (cpuset_getaffinity(level, which, id, sizeof(mask), &mask) != 0)
175 err(EXIT_FAILURE, "getaffinity");
176 printf("%s %jd%s mask: ", whichnames[which], (intmax_t)id,
177 levelnames[level]);
178 printset(&mask);
179 exit(EXIT_SUCCESS);
180 }
181
182 static void
printsetid(void)183 printsetid(void)
184 {
185 cpusetid_t setid;
186
187 /*
188 * Only LEVEL_WHICH && WHICH_CPUSET has a numbered id.
189 */
190 if (level == CPU_LEVEL_WHICH && !sflag)
191 level = CPU_LEVEL_CPUSET;
192 if (cpuset_getid(level, which, id, &setid))
193 err(errno, "getid");
194 printf("%s %jd%s id: %d\n", whichnames[which], (intmax_t)id,
195 levelnames[level], setid);
196 }
197
198 int
main(int argc,char * argv[])199 main(int argc, char *argv[])
200 {
201 cpusetid_t setid;
202 cpuset_t mask;
203 lwpid_t tid;
204 pid_t pid;
205 int ch;
206
207 CPU_ZERO(&mask);
208 level = CPU_LEVEL_WHICH;
209 which = CPU_WHICH_PID;
210 id = pid = tid = setid = -1;
211 while ((ch = getopt(argc, argv, "Ccd:gij:l:p:rs:t:x:")) != -1) {
212 switch (ch) {
213 case 'C':
214 Cflag = 1;
215 break;
216 case 'c':
217 cflag = 1;
218 level = CPU_LEVEL_CPUSET;
219 break;
220 case 'd':
221 dflag = 1;
222 which = CPU_WHICH_DOMAIN;
223 id = atoi(optarg);
224 break;
225 case 'g':
226 gflag = 1;
227 break;
228 case 'i':
229 iflag = 1;
230 break;
231 case 'j':
232 jflag = 1;
233 which = CPU_WHICH_JAIL;
234 id = atoi(optarg);
235 break;
236 case 'l':
237 lflag = 1;
238 parselist(optarg, &mask);
239 break;
240 case 'p':
241 pflag = 1;
242 which = CPU_WHICH_PID;
243 id = pid = atoi(optarg);
244 break;
245 case 'r':
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 (dflag + jflag + xflag + sflag + pflag + tflag > 1)
275 usage();
276 if (iflag)
277 printsetid();
278 else
279 printaffinity();
280 exit(EXIT_SUCCESS);
281 }
282 if (dflag || iflag || rflag)
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 || 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 && (jflag || !pflag || sflag || tflag || xflag))
310 usage();
311 if (!lflag && cflag)
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
349 static void
usage(void)350 usage(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 [-c] [-l cpu-list] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
361 fprintf(stderr,
362 " cpuset -g [-cir] [-d domain | -j jailid | -p pid | -t tid | -s setid |\n"
363 " -x irq]\n");
364 exit(1);
365 }
366