usr.bin/cpuset/cpuset.c

   1 /*-
   2  * SPDX-License-Identifier: BSD-2-Clause-FreeBSD
   3  *
   4  * Copyright (c) 2007, 2008     Jeffrey Roberson <jeff@freebsd.org>
   5  * All rights reserved.
   6  *
   7  * Copyright (c) 2008 Nokia Corporation
   8  * All rights reserved.
   9  *
  10  * Redistribution and use in source and binary forms, with or without
  11  * modification, are permitted provided that the following conditions
  12  * are met:
  13  * 1. Redistributions of source code must retain the above copyright
  14  *    notice, this list of conditions and the following disclaimer.
  15  * 2. Redistributions in binary form must reproduce the above copyright
  16  *    notice, this list of conditions and the following disclaimer in the
  17  *    documentation and/or other materials provided with the distribution.
  18  *
  19  * THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
  20  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  21  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  22  * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
  23  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  24  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  25  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  26  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  27  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  28  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  29  * SUCH DAMAGE.
  30  */
  31
  32 #include <sys/cdefs.h>
  33 __FBSDID("$FreeBSD$");
  34
  35 #include <sys/param.h>
  36 #include <sys/types.h>
  37 #include <sys/time.h>
  38 #include <sys/resource.h>
  39 #include <sys/cpuset.h>
  40
  41 #include <ctype.h>
  42 #include <err.h>
  43 #include <errno.h>
  44 #include <limits.h>
  45 #include <stdio.h>
  46 #include <stdlib.h>
  47 #include <stdint.h>
  48 #include <unistd.h>
  49 #include <string.h>
  50
  51 static int Cflag;
  52 static int cflag;
  53 static int dflag;
  54 static int gflag;
  55 static int iflag;
  56 static int jflag;
  57 static int lflag;
  58 static int pflag;
  59 static int rflag;
  60 static int sflag;
  61 static int tflag;
  62 static int xflag;
  63 static id_t id;
  64 static cpulevel_t level;
  65 static cpuwhich_t which;
  66
  67 static void usage(void);
  68
  69 static void printset(cpuset_t *mask);
  70
  71 static void
  72 parselist(char *list, cpuset_t *mask)
  73 {
  74         enum { NONE, NUM, DASH } state;
  75         int lastnum;
  76         int curnum;
  77         char *l;
  78
  79         if (strcasecmp(list, "all") == 0) {
  80                 if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
  81                     sizeof(*mask), mask) != 0)
  82                         err(EXIT_FAILURE, "getaffinity");
  83                 return;
  84         }
  85         state = NONE;
  86         curnum = lastnum = 0;
  87         for (l = list; *l != '\0';) {
  88                 if (isdigit(*l)) {
  89                         curnum = atoi(l);
  90                         if (curnum > CPU_SETSIZE)
  91                                 errx(EXIT_FAILURE,
  92                                     "Only %d cpus supported", CPU_SETSIZE);
  93                         while (isdigit(*l))
  94                                 l++;
  95                         switch (state) {
  96                         case NONE:
  97                                 lastnum = curnum;
  98                                 state = NUM;
  99                                 break;
 100                         case DASH:
 101                                 for (; lastnum <= curnum; lastnum++)
 102                                         CPU_SET(lastnum, mask);
 103                                 state = NONE;
 104                                 break;
 105                         case NUM:
 106                         default:
 107                                 goto parserr;
 108                         }
 109                         continue;
 110                 }
 111                 switch (*l) {
 112                 case ',':
 113                         switch (state) {
 114                         case NONE:
 115                                 break;
 116                         case NUM:
 117                                 CPU_SET(curnum, mask);
 118                                 state = NONE;
 119                                 break;
 120                         case DASH:
 121                                 goto parserr;
 122                                 break;
 123                         }
 124                         break;
 125                 case '-':
 126                         if (state != NUM)
 127                                 goto parserr;
 128                         state = DASH;
 129                         break;
 130                 default:
 131                         goto parserr;
 132                 }
 133                 l++;
 134         }
 135         switch (state) {
 136                 case NONE:
 137                         break;
 138                 case NUM:
 139                         CPU_SET(curnum, mask);
 140                         break;
 141                 case DASH:
 142                         goto parserr;
 143         }
 144         return;
 145 parserr:
 146         errx(EXIT_FAILURE, "Malformed cpu-list %s", list);
 147 }
 148
 149 static void
 150 printset(cpuset_t *mask)
 151 {
 152         int once;
 153         int cpu;
 154
 155         for (once = 0, cpu = 0; cpu < CPU_SETSIZE; cpu++) {
 156                 if (CPU_ISSET(cpu, mask)) {
 157                         if (once == 0) {
 158                                 printf("%d", cpu);
 159                                 once = 1;
 160                         } else
 161                                 printf(", %d", cpu);
 162                 }
 163         }
 164         printf("\n");
 165 }
 166
 167 static const char *whichnames[] = { NULL, "tid", "pid", "cpuset", "irq", "jail",
 168                                     "domain" };
 169 static const char *levelnames[] = { NULL, " root", " cpuset", "" };
 170
 171 static void
 172 printaffinity(void)
 173 {
 174         cpuset_t mask;
 175
 176         if (cpuset_getaffinity(level, which, id, sizeof(mask), &mask) != 0)
 177                 err(EXIT_FAILURE, "getaffinity");
 178         printf("%s %jd%s mask: ", whichnames[which], (intmax_t)id,
 179             levelnames[level]);
 180         printset(&mask);
 181         exit(EXIT_SUCCESS);
 182 }
 183
 184 static void
 185 printsetid(void)
 186 {
 187         cpusetid_t setid;
 188
 189         /*
 190          * Only LEVEL_WHICH && WHICH_CPUSET has a numbered id.
 191          */
 192         if (level == CPU_LEVEL_WHICH && !sflag)
 193                 level = CPU_LEVEL_CPUSET;
 194         if (cpuset_getid(level, which, id, &setid))
 195                 err(errno, "getid");
 196         printf("%s %jd%s id: %d\n", whichnames[which], (intmax_t)id,
 197             levelnames[level], setid);
 198 }
 199
 200 int
 201 main(int argc, char *argv[])
 202 {
 203         cpusetid_t setid;
 204         cpuset_t mask;
 205         lwpid_t tid;
 206         pid_t pid;
 207         int ch;
 208
 209         CPU_ZERO(&mask);
 210         level = CPU_LEVEL_WHICH;
 211         which = CPU_WHICH_PID;
 212         id = pid = tid = setid = -1;
 213         while ((ch = getopt(argc, argv, "Ccd:gij:l:p:rs:t:x:")) != -1) {
 214                 switch (ch) {
 215                 case 'C':
 216                         Cflag = 1;
 217                         break;
 218                 case 'c':
 219                         cflag = 1;
 220                         level = CPU_LEVEL_CPUSET;
 221                         break;
 222                 case 'd':
 223                         dflag = 1;
 224                         which = CPU_WHICH_DOMAIN;
 225                         id = atoi(optarg);
 226                         break;
 227                 case 'g':
 228                         gflag = 1;
 229                         break;
 230                 case 'i':
 231                         iflag = 1;
 232                         break;
 233                 case 'j':
 234                         jflag = 1;
 235                         which = CPU_WHICH_JAIL;
 236                         id = atoi(optarg);
 237                         break;
 238                 case 'l':
 239                         lflag = 1;
 240                         parselist(optarg, &mask);
 241                         break;
 242                 case 'p':
 243                         pflag = 1;
 244                         which = CPU_WHICH_PID;
 245                         id = pid = atoi(optarg);
 246                         break;
 247                 case 'r':
 248                         level = CPU_LEVEL_ROOT;
 249                         rflag = 1;
 250                         break;
 251                 case 's':
 252                         sflag = 1;
 253                         which = CPU_WHICH_CPUSET;
 254                         id = setid = atoi(optarg);
 255                         break;
 256                 case 't':
 257                         tflag = 1;
 258                         which = CPU_WHICH_TID;
 259                         id = tid = atoi(optarg);
 260                         break;
 261                 case 'x':
 262                         xflag = 1;
 263                         which = CPU_WHICH_IRQ;
 264                         id = atoi(optarg);
 265                         break;
 266                 default:
 267                         usage();
 268                 }
 269         }
 270         argc -= optind;
 271         argv += optind;
 272         if (gflag) {
 273                 if (argc || Cflag || lflag)
 274                         usage();
 275                 /* Only one identity specifier. */
 276                 if (dflag + jflag + xflag + sflag + pflag + tflag > 1)
 277                         usage();
 278                 if (iflag)
 279                         printsetid();
 280                 else
 281                         printaffinity();
 282                 exit(EXIT_SUCCESS);
 283         }
 284         if (dflag || iflag || rflag)
 285                 usage();
 286         /*
 287          * The user wants to run a command with a set and possibly cpumask.
 288          */
 289         if (argc) {
 290                 if (Cflag || pflag || tflag || xflag || jflag)
 291                         usage();
 292                 if (sflag) {
 293                         if (cpuset_setid(CPU_WHICH_PID, -1, setid))
 294                                 err(argc, "setid");
 295                 } else {
 296                         if (cpuset(&setid))
 297                                 err(argc, "newid");
 298                 }
 299                 if (lflag) {
 300                         if (cpuset_setaffinity(level, CPU_WHICH_PID,
 301                             -1, sizeof(mask), &mask) != 0)
 302                                 err(EXIT_FAILURE, "setaffinity");
 303                 }
 304                 errno = 0;
 305                 execvp(*argv, argv);
 306                 err(errno == ENOENT ? 127 : 126, "%s", *argv);
 307         }
 308         /*
 309          * We're modifying something that presently exists.
 310          */
 311         if (Cflag && (jflag || !pflag || sflag || tflag || xflag))
 312                 usage();
 313         if (!lflag && cflag)
 314                 usage();
 315         if (!lflag && !(Cflag || sflag))
 316                 usage();
 317         /* You can only set a mask on a thread. */
 318         if (tflag && (sflag | pflag | xflag | jflag))
 319                 usage();
 320         /* You can only set a mask on an irq. */
 321         if (xflag && (jflag | pflag | sflag | tflag))
 322                 usage();
 323         if (Cflag) {
 324                 /*
 325                  * Create a new cpuset and move the specified process
 326                  * into the set.
 327                  */
 328                 if (cpuset(&setid) < 0)
 329                         err(EXIT_FAILURE, "newid");
 330                 sflag = 1;
 331         }
 332         if (pflag && sflag) {
 333                 if (cpuset_setid(CPU_WHICH_PID, pid, setid))
 334                         err(EXIT_FAILURE, "setid");
 335                 /*
 336                  * If the user specifies a set and a list we want the mask
 337                  * to effect the pid and not the set.
 338                  */
 339                 which = CPU_WHICH_PID;
 340                 id = pid;
 341         }
 342         if (lflag) {
 343                 if (cpuset_setaffinity(level, which, id, sizeof(mask),
 344                     &mask) != 0)
 345                         err(EXIT_FAILURE, "setaffinity");
 346         }
 347
 348         exit(EXIT_SUCCESS);
 349 }
 350
 351 static void
 352 usage(void)
 353 {
 354
 355         fprintf(stderr,
 356             "usage: cpuset [-l cpu-list] [-s setid] cmd ...\n");
 357         fprintf(stderr,
 358             "       cpuset [-l cpu-list] [-s setid] -p pid\n");
 359         fprintf(stderr,
 360             "       cpuset [-c] [-l cpu-list] -C -p pid\n");
 361         fprintf(stderr,
 362             "       cpuset [-c] [-l cpu-list] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
 363         fprintf(stderr,
 364             "       cpuset -g [-cir] [-d domain | -j jailid | -p pid | -t tid | -s setid |\n"
 365             "              -x irq]\n");
 366         exit(1);
 367 }