1/*	$NetBSD: schedctl.c,v 1.14 2009/05/16 23:59:56 rmind Exp $	*/
2
3/*
4 * Copyright (c) 2008, Mindaugas Rasiukevicius <rmind at NetBSD org>
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 AUTHOR AND CONTRIBUTORS ``AS IS'' AND
17 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
20 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
21 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
22 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
23 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
24 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
25 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
26 * SUCH DAMAGE.
27 */
28
29/*
30 * schedctl(8) - a program to control scheduling of processes and threads.
31 */
32
33#include <sys/cdefs.h>
34
35#ifndef lint
36__RCSID("$NetBSD: schedctl.c,v 1.14 2009/05/16 23:59:56 rmind Exp $");
37#endif
38
39#include <stdio.h>
40#include <stdlib.h>
41#include <string.h>
42
43#include <err.h>
44#include <fcntl.h>
45#include <kvm.h>
46#include <unistd.h>
47
48#include <sys/pset.h>
49#include <sys/sched.h>
50#include <sys/sysctl.h>
51#include <sys/types.h>
52
53static const char *class_str[] = {
54	"SCHED_OTHER",
55	"SCHED_FIFO",
56	"SCHED_RR",
57	NULL
58};
59
60static void	sched_set(pid_t, lwpid_t, int, struct sched_param *, cpuset_t *);
61static void	thread_info(pid_t, lwpid_t);
62static cpuset_t	*makecpuset(char *);
63static char	*showcpuset(cpuset_t *);
64__dead static void	usage(void);
65
66static u_int	ncpu;
67
68int
69main(int argc, char **argv)
70{
71	kvm_t *kd;
72	struct kinfo_lwp *lwp_list, *lwp;
73	struct sched_param *sp;
74	cpuset_t *cpuset;
75	int i, count, ch, policy;
76	pid_t pid;
77	lwpid_t lid;
78	bool set;
79
80	ncpu = sysconf(_SC_NPROCESSORS_CONF);
81
82	pid = lid = 0;
83	cpuset = NULL;
84	set = false;
85
86	sp = calloc(1, sizeof(struct sched_param));
87	if (sp == NULL)
88		err(EXIT_FAILURE, "calloc");
89
90	sp->sched_priority = PRI_NONE;
91	policy = SCHED_NONE;
92
93	while ((ch = getopt(argc, argv, "A:C:P:p:t:")) != -1) {
94		switch (ch) {
95		case 'p':
96			/* PID */
97			pid = atoi(optarg);
98			break;
99		case 't':
100			/* Thread (LWP) ID */
101			lid = atoi(optarg);
102			break;
103		case 'A':
104			/* Affinity */
105			cpuset = makecpuset(optarg);
106			if (cpuset == NULL) {
107				fprintf(stderr, "%s: invalid CPU value\n",
108				    getprogname());
109				exit(EXIT_FAILURE);
110			}
111			break;
112		case 'C':
113			/* Scheduling class */
114			for (policy = 0; class_str[policy] != NULL; policy++) {
115				if (strcasecmp(optarg, class_str[policy]) == 0)
116					break;
117			}
118			if (class_str[policy] == NULL)
119				policy = atoi(optarg);
120			if (policy < SCHED_OTHER || policy > SCHED_RR) {
121				fprintf(stderr,
122				    "%s: invalid scheduling class\n",
123				    getprogname());
124				exit(EXIT_FAILURE);
125			}
126			set = true;
127			break;
128		case 'P':
129			/* Priority */
130			sp->sched_priority = atoi(optarg);
131			if (sp->sched_priority < sysconf(_SC_SCHED_PRI_MIN) ||
132			    sp->sched_priority > sysconf(_SC_SCHED_PRI_MAX)) {
133				fprintf(stderr, "%s: invalid priority\n",
134				    getprogname());
135				exit(EXIT_FAILURE);
136			}
137			set = true;
138			break;
139		default:
140			usage();
141		}
142	}
143
144	/* At least PID must be specified */
145	if (pid == 0) {
146		if (argv[optind] == NULL || lid != 0)
147			usage();
148		pid = getpid();
149	} else {
150		if (argv[optind] != NULL)
151			usage();
152	}
153
154	/* Set the scheduling information for thread/process */
155	sched_set(pid, lid, policy, set ? sp : NULL, cpuset);
156
157	/* Show information about each thread */
158	if (pid != getpid()) {
159		kd = kvm_open(NULL, NULL, NULL, KVM_NO_FILES, "kvm_open");
160		if (kd == NULL)
161			err(EXIT_FAILURE, "kvm_open");
162		lwp_list = kvm_getlwps(kd, pid, 0, sizeof(struct kinfo_lwp), &count);
163		if (lwp_list == NULL)
164			err(EXIT_FAILURE, "kvm_getlwps");
165		for (lwp = lwp_list, i = 0; i < count; lwp++, i++) {
166			if (lid && lid != lwp->l_lid)
167				continue;
168			if (lwp->l_stat == LSIDL || lwp->l_stat == LSZOMB)
169				continue;
170			thread_info(pid, lwp->l_lid);
171		}
172		kvm_close(kd);
173		free(sp);
174		cpuset_destroy(cpuset);
175		return 0;
176	}
177
178	(void)execvp(argv[optind], argv + optind);
179	err(EXIT_FAILURE, "execvp");
180}
181
182static void
183sched_set(pid_t pid, lwpid_t lid, int policy,
184    struct sched_param *sp, cpuset_t *cpuset)
185{
186	int error;
187
188	if (sp) {
189		/* Set the scheduling parameters for the thread */
190		error = _sched_setparam(pid, lid, policy, sp);
191		if (error < 0)
192			err(EXIT_FAILURE, "_sched_setparam");
193	}
194	if (cpuset) {
195		/* Set the CPU-set for affinity */
196		error = _sched_setaffinity(pid, lid,
197		    cpuset_size(cpuset), cpuset);
198		if (error < 0)
199			err(EXIT_FAILURE, "_sched_setaffinity");
200	}
201}
202
203static void
204thread_info(pid_t pid, lwpid_t lid)
205{
206	struct sched_param sp;
207	cpuset_t *cpuset;
208	char *cpus;
209	int error, policy;
210
211	cpuset = cpuset_create();
212	if (cpuset == NULL)
213		err(EXIT_FAILURE, "cpuset_create");
214
215	error = _sched_getparam(pid, lid, &policy, &sp);
216	if (error < 0)
217		err(EXIT_FAILURE, "_sched_getparam");
218
219	error = _sched_getaffinity(pid, lid, cpuset_size(cpuset), cpuset);
220	if (error < 0)
221		err(EXIT_FAILURE, "_sched_getaffinity");
222
223	printf("  LID:              %d\n", lid);
224	printf("  Priority:         %d\n", sp.sched_priority);
225	printf("  Class:            %s\n", class_str[policy]);
226
227	cpus = showcpuset(cpuset);
228	printf("  Affinity (CPUs):  %s\n", cpus);
229	free(cpus);
230
231	cpuset_destroy(cpuset);
232}
233
234static cpuset_t *
235makecpuset(char *str)
236{
237	cpuset_t *cpuset;
238	char *cpustr, *s;
239
240	if (str == NULL)
241		return NULL;
242
243	cpuset = cpuset_create();
244	if (cpuset == NULL)
245		err(EXIT_FAILURE, "cpuset_create");
246	cpuset_zero(cpuset);
247
248	cpustr = strdup(str);
249	if (cpustr == NULL)
250		err(EXIT_FAILURE, "strdup");
251	s = cpustr;
252
253	while (s != NULL) {
254		char *p;
255		int i;
256
257		/* Get the CPU number and validate the range */
258		p = strsep(&s, ",");
259		if (p == NULL) {
260			cpuset_destroy(cpuset);
261			cpuset = NULL;
262			break;
263		}
264		i = atoi(p);
265		if (i == -1) {
266			cpuset_zero(cpuset);
267			break;
268		}
269		if ((unsigned int)i >= ncpu) {
270			cpuset_destroy(cpuset);
271			cpuset = NULL;
272			break;
273		}
274
275		/* Set the bit */
276		cpuset_set(i, cpuset);
277	}
278
279	free(cpustr);
280	return cpuset;
281}
282
283static char *
284showcpuset(cpuset_t *cpuset)
285{
286	char *buf;
287	size_t size;
288	unsigned int i;
289
290	size = 3 * ncpu;	/* XXX */
291	buf = malloc(size + 1);
292	if (buf == NULL)
293		err(EXIT_FAILURE, "malloc");
294	memset(buf, '\0', size + 1);
295
296	for (i = 0; i < ncpu; i++)
297		if (cpuset_isset(i, cpuset))
298			snprintf(buf, size, "%s%d,", buf, i);
299
300	i = strlen(buf);
301	if (i != 0) {
302		buf[i - 1] = '\0';
303	} else {
304		strncpy(buf, "<none>", size);
305	}
306
307	return buf;
308}
309
310static void
311usage(void)
312{
313
314	fprintf(stderr, "usage: %s [-A processor] [-C class] "
315	    "[-P priority] [-t lid] {-p pid|command}\n", getprogname());
316	exit(EXIT_FAILURE);
317}
318