1/*	$NetBSD: psrset.c,v 1.2 2009/01/18 10:04:15 lukem Exp $	*/
2
3/*-
4 * Copyright (c) 2008 The NetBSD Foundation, Inc.
5 * All rights reserved.
6 *
7 * Redistribution and use in source and binary forms, with or without
8 * modification, are permitted provided that the following conditions
9 * are met:
10 * 1. Redistributions of source code must retain the above copyright
11 *    notice, this list of conditions and the following disclaimer.
12 * 2. Redistributions in binary form must reproduce the above copyright
13 *    notice, this list of conditions and the following disclaimer in the
14 *    documentation and/or other materials provided with the distribution.
15 *
16 * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
17 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
18 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
19 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
20 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
21 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
22 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
23 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
24 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
25 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
26 * POSSIBILITY OF SUCH DAMAGE.
27 */
28
29#include <sys/cdefs.h>
30#ifndef lint
31__RCSID("$NetBSD: psrset.c,v 1.2 2009/01/18 10:04:15 lukem Exp $");
32#endif
33
34#include <sys/types.h>
35#include <sys/pset.h>
36#include <sys/sched.h>
37#include <sys/sysctl.h>
38
39#include <stdio.h>
40#include <stdlib.h>
41#include <string.h>
42#include <fcntl.h>
43#include <err.h>
44#include <unistd.h>
45#include <ctype.h>
46
47__dead static void	usage(void);
48static int	eatopt(char **);
49static int	cmd_a(char **, int);
50static int	cmd_b(char **, int);
51static int	cmd_c(char **, int);
52static int	cmd_d(char **, int);
53static int	cmd_e(char **, int);
54static int	cmd_i(char **, int);
55static int	cmd_p(char **, int);
56static int	cmd_r(char **, int);
57static int	cmd_u(char **, int);
58static int	cmd__(char **, int);
59
60static psetid_t	psid;
61static int	ncpu;
62static cpuset_t	*cpuset;
63
64int
65main(int argc, char **argv)
66{
67	int (*cmd)(char **, int);
68	int off;
69
70	ncpu = sysconf(_SC_NPROCESSORS_CONF);
71	cpuset = cpuset_create();
72	if (cpuset == NULL)
73		err(EXIT_FAILURE, "cpuset_create");
74	cpuset_zero(cpuset);
75
76	if (argc == 1 || argv[1][0] != '-') {
77		cmd = cmd_i;
78		off = 1;
79	} else if (strncmp(argv[1], "-a", 2) == 0) {
80		cmd = cmd_a;
81		off = eatopt(argv);
82	} else if (strncmp(argv[1], "-b", 2) == 0) {
83		cmd = cmd_b;
84		off = eatopt(argv);
85	} else if (strncmp(argv[1], "-c", 2) == 0) {
86		cmd = cmd_c;
87		off = eatopt(argv);
88	} else if (strncmp(argv[1], "-d", 2) == 0) {
89		cmd = cmd_d;
90		off = eatopt(argv);
91	} else if (strncmp(argv[1], "-e", 2) == 0) {
92		cmd = cmd_e;
93		off = eatopt(argv);
94	} else if (strncmp(argv[1], "-i", 2) == 0) {
95		cmd = cmd_i;
96		off = eatopt(argv);
97	} else if (strncmp(argv[1], "-p", 2) == 0) {
98		cmd = cmd_p;
99		off = eatopt(argv);
100	} else if (strncmp(argv[1], "-r", 2) == 0) {
101		cmd = cmd_r;
102		off = eatopt(argv);
103	} else if (strncmp(argv[1], "-u", 2) == 0) {
104		cmd = cmd_u;
105		off = eatopt(argv);
106	} else {
107		cmd = cmd__;
108		off = 0;
109	}
110
111	return (*cmd)(argv + off, argc - off);
112}
113
114static int
115eatopt(char **argv)
116{
117
118	if (argv[1][2] != '\0') {
119		argv[1] += 2;
120		return 1;
121	}
122	return 2;
123}
124
125static int
126getint(char *p)
127{
128	char *q;
129	int rv;
130
131	rv = (int)strtol(p, &q, 10);
132	if (q == p || *q != '\0')
133		usage();
134	return rv;
135}
136
137static void
138usage(void)
139{
140
141	fprintf(stderr, "usage:\n"
142	    "\tpsrset [setid ...]\n"
143	    "\tpsrset -a setid cpuid ...\n"
144	    "\tpsrset -b setid pid ...\n"
145	    "\tpsrset -c [cpuid ...]\n"
146	    "\tpsrset -d setid\n"
147	    "\tpsrset -e setid command\n"
148	    "\tpsrset -i [setid ...]\n"
149	    "\tpsrset -p\n"
150	    "\tpsrset -r cpuid ...\n"
151	    "\tpsrset -u pid ...\n");
152
153	exit(EXIT_FAILURE);
154}
155
156static void
157makecpuset(char **argv)
158{
159	char *p, *q;
160	int i, j;
161
162	if (*argv == NULL) {
163		for (i = 0; i < ncpu; i++)
164			cpuset_set(i, cpuset);
165		return;
166	}
167
168	for (; *argv != NULL; argv++) {
169		if (!isdigit((unsigned)**argv))
170			usage();
171		i = (int)strtol(*argv, &p, 10);
172		if ((*p != '\0' && *p != '-') || p == *argv)
173			usage();
174		if (*p == '-')  {
175			if (!isdigit((unsigned)p[1]))
176				usage();
177			j = (int)strtol(p + 1, &q, 10);
178			if (q == p || *q != '\0')
179				usage();
180		} else {
181			j = i;
182		}
183		if (i >= ncpu) {
184			errx(EXIT_FAILURE, "value out of range");
185		}
186		while (i <= j)
187			cpuset_set(i++, cpuset);
188	}
189}
190
191/*
192 * Assign processors to set.
193 */
194static int
195cmd_a(char **argv, int argc)
196{
197	int i;
198
199	if (argc < 2)
200		usage();
201	psid = getint(argv[0]);
202	makecpuset(argv + 1);
203	for (i = 0; i < ncpu; i++) {
204		if (!cpuset_isset(i, cpuset))
205			continue;
206		if (pset_assign(psid, i, NULL))
207			err(EXIT_FAILURE, "pset_assign");
208	}
209
210	return 0;
211}
212
213/*
214 * Bind LWPs within processes to set.
215 */
216static int
217cmd_b(char **argv, int argc)
218{
219
220	if (argc < 2)
221		usage();
222	psid = getint(*argv);
223	for (argv++; *argv != NULL; argv++) {
224		if (pset_bind(psid, P_PID, (idtype_t)getint(*argv), NULL))
225			err(EXIT_FAILURE, "pset_bind");
226	}
227
228	return 0;
229}
230
231/*
232 * Create set.
233 */
234static int
235cmd_c(char **argv, int argc)
236{
237	int i;
238
239	if (pset_create(&psid))
240		err(EXIT_FAILURE, "pset_create");
241	printf("%d\n", (int)psid);
242	if (argc != 0)
243		makecpuset(argv);
244	for (i = 0; i < ncpu; i++) {
245		if (!cpuset_isset(i, cpuset))
246			continue;
247		if (pset_assign(psid, i, NULL))
248			err(EXIT_FAILURE, "pset_assign");
249	}
250
251	return 0;
252}
253
254/*
255 * Destroy set.
256 */
257static int
258cmd_d(char **argv, int argc)
259{
260
261	if (argc != 1)
262		usage();
263	if (pset_destroy(getint(argv[0])))
264		err(EXIT_FAILURE, "pset_destroy");
265
266	return 0;
267}
268
269/*
270 * Execute command in set.
271 */
272static int
273cmd_e(char **argv, int argc)
274{
275
276	if (argc < 2)
277		usage();
278	if (pset_bind(getint(argv[0]), P_PID, getpid(), NULL))
279		err(EXIT_FAILURE, "pset_bind");
280	(void)execvp(argv[1], argv + 1);
281	return 1;
282}
283
284/*
285 * Print info about each set.
286 */
287static int
288cmd_i(char **argv, int argc)
289{
290	char buf[1024];
291	size_t len;
292	char *p, *q;
293	int i, j, k;
294
295	len = sizeof(buf);
296	if (sysctlbyname("kern.pset.list", buf, &len,  NULL, 0) == -1)
297		err(EXIT_FAILURE, "kern.pset.list");
298
299	p = buf;
300	while ((q = strsep(&p, ",")) != NULL) {
301		i = (int)strtol(q, &q, 10);
302		if (*q != ':')
303			errx(EXIT_FAILURE, "bad kern.pset.list");
304		printf("%s processor set %d: ",
305		    (atoi(q + 1) == 1 ? "system" : "user"), i);
306		for (j = 0, k = 0; j < ncpu; j++) {
307			if (pset_assign(PS_QUERY, j, &psid))
308				err(EXIT_FAILURE, "pset_assign");
309			if (psid == i) {
310				if (k++ == 0)
311					printf("processor(s)");
312				printf(" %d", j);
313			}
314		}
315		if (k == 0)
316			printf("empty");
317		putchar('\n');
318	}
319
320	return 0;
321}
322
323/*
324 * Print set ID for each processor.
325 */
326static int
327cmd_p(char **argv, int argc)
328{
329	int i;
330
331	makecpuset(argv);
332
333	for (i = 0; i < ncpu; i++) {
334		if (!cpuset_isset(i, cpuset))
335			continue;
336		if (pset_assign(PS_QUERY, i, &psid))
337			err(EXIT_FAILURE, "ioctl");
338		printf("processor %d: ", i);
339		if (psid == PS_NONE)
340			printf("not assigned\n");
341		else
342			printf("%d\n", (int)psid);
343	}
344
345	return 0;
346}
347
348/*
349 * Remove CPU from set and return to system.
350 */
351static int
352cmd_r(char **argv, int argc)
353{
354	int i;
355
356	makecpuset(argv);
357
358	for (i = 0; i < ncpu; i++) {
359		if (!cpuset_isset(i, cpuset))
360			continue;
361		if (pset_assign(PS_NONE, i, NULL))
362			err(EXIT_FAILURE, "ioctl");
363	}
364
365	return 0;
366}
367
368/*
369 * Unbind LWPs within process.
370 */
371static int
372cmd_u(char **argv, int argc)
373{
374
375	if (argc < 1)
376		usage();
377	for (; *argv != NULL; argv++) {
378		if (pset_bind(PS_NONE, P_PID, (idtype_t)getint(*argv), NULL))
379			err(EXIT_FAILURE, "pset_bind");
380	}
381
382	return 0;
383}
384
385
386/*
387 * Dummy.
388 */
389static int
390cmd__(char **argv, int argc)
391{
392
393	usage();
394	return 0;
395}
396