usr.sbin/psrset/psrset.c

   1 /*      $NetBSD: psrset.c,v 1.1 2008/06/22 13:53:59 ad 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.1 2008/06/22 13:53:59 ad 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 static void     usage(void);
  48 static int      eatopt(char **);
  49 static int      cmd_a(char **, int);
  50 static int      cmd_b(char **, int);
  51 static int      cmd_c(char **, int);
  52 static int      cmd_d(char **, int);
  53 static int      cmd_e(char **, int);
  54 static int      cmd_i(char **, int);
  55 static int      cmd_p(char **, int);
  56 static int      cmd_r(char **, int);
  57 static int      cmd_u(char **, int);
  58 static int      cmd__(char **, int);
  59
  60 static psetid_t psid;
  61 static int      ncpu;
  62 static cpuset_t *cpuset;
  63
  64 int
  65 main(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
 114 static int
 115 eatopt(char **argv)
 116 {
 117
 118         if (argv[1][2] != '\0') {
 119                 argv[1] += 2;
 120                 return 1;
 121         }
 122         return 2;
 123 }
 124
 125 static int
 126 getint(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
 137 static void
 138 usage(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
 156 static void
 157 makecpuset(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  */
 194 static int
 195 cmd_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  */
 216 static int
 217 cmd_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  */
 234 static int
 235 cmd_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  */
 257 static int
 258 cmd_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  */
 272 static int
 273 cmd_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  */
 287 static int
 288 cmd_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  */
 326 static int
 327 cmd_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  */
 351 static int
 352 cmd_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  */
 371 static int
 372 cmd_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  */
 389 static int
 390 cmd__(char **argv, int argc)
 391 {
 392
 393         usage();
 394         return 0;
 395 }