sys: Extract CPUMASK macros to new <machine/cpumask.h>
[dragonfly.git] / sbin / usched / usched.c
1 /*
2  * Copyright (c) 2012 The DragonFly Project.  All rights reserved.
3  *
4  * This code is derived from software contributed to The DragonFly Project
5  * by Matthew Dillon <dillon@backplane.com> and Thomas Nikolajsen
6  * <thomas.nikolajsen@mail.dk>
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  *
12  * 1. Redistributions of source code must retain the above copyright
13  *    notice, this list of conditions and the following disclaimer.
14  * 2. Redistributions in binary form must reproduce the above copyright
15  *    notice, this list of conditions and the following disclaimer in
16  *    the documentation and/or other materials provided with the
17  *    distribution.
18  * 3. Neither the name of The DragonFly Project nor the names of its
19  *    contributors may be used to endorse or promote products derived
20  *    from this software without specific, prior written permission.
21  *
22  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
23  * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
24  * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
25  * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
26  * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
27  * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
28  * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
29  * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
30  * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
31  * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
32  * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
33  * SUCH DAMAGE.
34  */
35
36 #include <sys/types.h>
37 #include <sys/usched.h>
38 #include <machine/cpumask.h>
39 #include <stdio.h>
40 #include <stdlib.h>
41 #include <unistd.h>
42 #include <string.h>
43
44 static void usage(void);
45
46 int DebugOpt;
47
48 int
49 main(int ac, char **av)
50 {
51         int ch;
52         int res;
53         char *sched = NULL;
54         char *cpustr = NULL;
55         char *sched_cpustr = NULL;
56         char *p = NULL;
57         cpumask_t cpumask;
58         int cpuid;
59         pid_t pid = getpid();  /* See usched_set(2) - BUGS */
60
61         CPUMASK_ASSZERO(cpumask);
62
63         while ((ch = getopt(ac, av, "d")) != -1) {
64                 switch (ch) {
65                 case 'd':
66                         DebugOpt = 1;
67                         break;
68                 default:
69                         usage();
70                         /* NOTREACHED */
71                 }
72         }
73         ac -= optind;
74         av += optind;
75
76         if (ac < 2) {
77                 usage();
78                 /* NOTREACHED */
79         }
80         sched_cpustr = strdup(av[0]);
81         sched = strsep(&sched_cpustr, ":");
82         if (strcmp(sched, "default") == 0)
83                 fprintf(stderr, "Ignoring scheduler == \"default\": not implemented\n");
84         cpustr = strsep(&sched_cpustr, "");
85         if (strlen(sched) == 0 && cpustr == NULL) {
86                 usage();
87                 /* NOTREACHED */
88         }
89
90         /*
91          * XXX needs expanded support for > 64 cpus
92          */
93         if (cpustr != NULL) {
94                 uint64_t v;
95
96                 v = (uint64_t)strtoull(cpustr, NULL, 0);
97                 for (cpuid = 0; cpuid < (int)sizeof(v) * 8; ++cpuid) {
98                         if (v & (1LU << cpuid))
99                                 CPUMASK_ORBIT(cpumask, cpuid);
100                 }
101         }
102
103         if (strlen(sched) != 0) {
104                 if (DebugOpt)
105                         fprintf(stderr, "DEBUG: USCHED_SET_SCHEDULER: scheduler: %s\n", sched);
106                 res = usched_set(pid, USCHED_SET_SCHEDULER, sched, strlen(sched));
107                 if (res != 0) {
108                         asprintf(&p, "usched_set(%d, USCHED_SET_SCHEDULER, \"%s\", %d)",
109                                 pid, sched, (int)strlen(sched));
110                         perror(p);
111                         exit(1);
112                 }
113         }
114         if (CPUMASK_TESTNZERO(cpumask)) {
115                 for (cpuid = 0; cpuid < (int)sizeof(cpumask) * 8; ++cpuid) {
116                         if (CPUMASK_TESTBIT(cpumask, cpuid))
117                                 break;
118                 }
119                 if (DebugOpt) {
120                         fprintf(stderr, "DEBUG: USCHED_SET_CPU: cpuid: %d\n",
121                                 cpuid);
122                 }
123                 res = usched_set(pid, USCHED_SET_CPU, &cpuid, sizeof(int));
124                 if (res != 0) {
125                         asprintf(&p, "usched_set(%d, USCHED_SET_CPU, &%d, %d)",
126                                 pid, cpuid, (int)sizeof(int));
127                         perror(p);
128                         exit(1);
129                 }
130                 CPUMASK_NANDBIT(cpumask, cpuid);
131                 while (CPUMASK_TESTNZERO(cpumask)) {
132                         ++cpuid;
133                         if (CPUMASK_TESTBIT(cpumask, cpuid) == 0)
134                                 continue;
135                         CPUMASK_NANDBIT(cpumask, cpuid);
136                         if (DebugOpt) {
137                                 fprintf(stderr,
138                                         "DEBUG: USCHED_ADD_CPU: cpuid: %d\n",
139                                         cpuid);
140                         }
141                         res = usched_set(pid, USCHED_ADD_CPU, &cpuid, sizeof(int));
142                         if (res != 0) {
143                                 asprintf(&p, "usched_set(%d, USCHED_ADD_CPU, &%d, %d)",
144                                         pid, cpuid, (int)sizeof(int));
145                                 perror(p);
146                                 exit(1);
147                         }
148                 }
149         }
150         execvp(av[1], av + 1);
151         exit(1);
152 }
153
154 static
155 void
156 usage(void)
157 {
158         fprintf(stderr,
159                 "usage: usched [-d] {scheduler[:cpumask] | :cpumask} "
160                 "program [argument ...]\n");
161         exit(1);
162 }