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 <jail.h>
  45 #include <limits.h>
  46 #include <stdio.h>
  47 #include <stdlib.h>
  48 #include <stdint.h>
  49 #include <unistd.h>
  50 #include <string.h>
  51
  52 static int Cflag;
  53 static int cflag;
  54 static int dflag;
  55 static int gflag;
  56 static int iflag;
  57 static int jflag;
  58 static int lflag;
  59 static int pflag;
  60 static int rflag;
  61 static int sflag;
  62 static int tflag;
  63 static int xflag;
  64 static id_t id;
  65 static cpulevel_t level;
  66 static cpuwhich_t which;
  67
  68 static void usage(void);
  69
  70 static void printset(cpuset_t *mask);
  71
  72 static void
  73 parselist(char *list, cpuset_t *mask)
  74 {
  75         enum { NONE, NUM, DASH } state;
  76         int lastnum;
  77         int curnum;
  78         char *l;
  79
  80         if (strcasecmp(list, "all") == 0) {
  81                 if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
  82                     sizeof(*mask), mask) != 0)
  83                         err(EXIT_FAILURE, "getaffinity");
  84                 return;
  85         }
  86         state = NONE;
  87         curnum = lastnum = 0;
  88         for (l = list; *l != '\0';) {
  89                 if (isdigit(*l)) {
  90                         curnum = atoi(l);
  91                         if (curnum > CPU_SETSIZE)
  92                                 errx(EXIT_FAILURE,
  93                                     "Only %d cpus supported", CPU_SETSIZE);
  94                         while (isdigit(*l))
  95                                 l++;
  96                         switch (state) {
  97                         case NONE:
  98                                 lastnum = curnum;
  99                                 state = NUM;
 100                                 break;
 101                         case DASH:
 102                                 for (; lastnum <= curnum; lastnum++)
 103                                         CPU_SET(lastnum, mask);
 104                                 state = NONE;
 105                                 break;
 106                         case NUM:
 107                         default:
 108                                 goto parserr;
 109                         }
 110                         continue;
 111                 }
 112                 switch (*l) {
 113                 case ',':
 114                         switch (state) {
 115                         case NONE:
 116                                 break;
 117                         case NUM:
 118                                 CPU_SET(curnum, mask);
 119                                 state = NONE;
 120                                 break;
 121                         case DASH:
 122                                 goto parserr;
 123                                 break;
 124                         }
 125                         break;
 126                 case '-':
 127                         if (state != NUM)
 128                                 goto parserr;
 129                         state = DASH;
 130                         break;
 131                 default:
 132                         goto parserr;
 133                 }
 134                 l++;
 135         }
 136         switch (state) {
 137                 case NONE:
 138                         break;
 139                 case NUM:
 140                         CPU_SET(curnum, mask);
 141                         break;
 142                 case DASH:
 143                         goto parserr;
 144         }
 145         return;
 146 parserr:
 147         errx(EXIT_FAILURE, "Malformed cpu-list %s", list);
 148 }
 149
 150 static void
 151 printset(cpuset_t *mask)
 152 {
 153         int once;
 154         int cpu;
 155
 156         for (once = 0, cpu = 0; cpu < CPU_SETSIZE; cpu++) {
 157                 if (CPU_ISSET(cpu, mask)) {
 158                         if (once == 0) {
 159                                 printf("%d", cpu);
 160                                 once = 1;
 161                         } else
 162                                 printf(", %d", cpu);
 163                 }
 164         }
 165         printf("\n");
 166 }
 167
 168 static const char *whichnames[] = { NULL, "tid", "pid", "cpuset", "irq", "jail",
 169                                     "domain" };
 170 static const char *levelnames[] = { NULL, " root", " cpuset", "" };
 171
 172 static void
 173 printaffinity(void)
 174 {
 175         cpuset_t mask;
 176
 177         if (cpuset_getaffinity(level, which, id, sizeof(mask), &mask) != 0)
 178                 err(EXIT_FAILURE, "getaffinity");
 179         printf("%s %jd%s mask: ", whichnames[which], (intmax_t)id,
 180             levelnames[level]);
 181         printset(&mask);
 182         exit(EXIT_SUCCESS);
 183 }
 184
 185 static void
 186 printsetid(void)
 187 {
 188         cpusetid_t setid;
 189
 190         /*
 191          * Only LEVEL_WHICH && WHICH_CPUSET has a numbered id.
 192          */
 193         if (level == CPU_LEVEL_WHICH && !sflag)
 194                 level = CPU_LEVEL_CPUSET;
 195         if (cpuset_getid(level, which, id, &setid))
 196                 err(errno, "getid");
 197         printf("%s %jd%s id: %d\n", whichnames[which], (intmax_t)id,
 198             levelnames[level], setid);
 199 }
 200
 201 int
 202 main(int argc, char *argv[])
 203 {
 204         cpusetid_t setid;
 205         cpuset_t mask;
 206         lwpid_t tid;
 207         pid_t pid;
 208         int ch;
 209
 210         CPU_ZERO(&mask);
 211         level = CPU_LEVEL_WHICH;
 212         which = CPU_WHICH_PID;
 213         id = pid = tid = setid = -1;
 214         while ((ch = getopt(argc, argv, "Ccd:gij:l:p:rs:t:x:")) != -1) {
 215                 switch (ch) {
 216                 case 'C':
 217                         Cflag = 1;
 218                         break;
 219                 case 'c':
 220                         cflag = 1;
 221                         level = CPU_LEVEL_CPUSET;
 222                         break;
 223                 case 'd':
 224                         dflag = 1;
 225                         which = CPU_WHICH_DOMAIN;
 226                         id = atoi(optarg);
 227                         break;
 228                 case 'g':
 229                         gflag = 1;
 230                         break;
 231                 case 'i':
 232                         iflag = 1;
 233                         break;
 234                 case 'j':
 235                         jflag = 1;
 236                         which = CPU_WHICH_JAIL;
 237                         id = jail_getid(optarg);
 238                         if (id < 0)
 239                                 errx(EXIT_FAILURE, "%s", jail_errmsg);
 240                         break;
 241                 case 'l':
 242                         lflag = 1;
 243                         parselist(optarg, &mask);
 244                         break;
 245                 case 'p':
 246                         pflag = 1;
 247                         which = CPU_WHICH_PID;
 248                         id = pid = atoi(optarg);
 249                         break;
 250                 case 'r':
 251                         level = CPU_LEVEL_ROOT;
 252                         rflag = 1;
 253                         break;
 254                 case 's':
 255                         sflag = 1;
 256                         which = CPU_WHICH_CPUSET;
 257                         id = setid = atoi(optarg);
 258                         break;
 259                 case 't':
 260                         tflag = 1;
 261                         which = CPU_WHICH_TID;
 262                         id = tid = atoi(optarg);
 263                         break;
 264                 case 'x':
 265                         xflag = 1;
 266                         which = CPU_WHICH_IRQ;
 267                         id = atoi(optarg);
 268                         break;
 269                 default:
 270                         usage();
 271                 }
 272         }
 273         argc -= optind;
 274         argv += optind;
 275         if (gflag) {
 276                 if (argc || Cflag || lflag)
 277                         usage();
 278                 /* Only one identity specifier. */
 279                 if (dflag + jflag + xflag + sflag + pflag + tflag > 1)
 280                         usage();
 281                 if (iflag)
 282                         printsetid();
 283                 else
 284                         printaffinity();
 285                 exit(EXIT_SUCCESS);
 286         }
 287         if (dflag || iflag || rflag)
 288                 usage();
 289         /*
 290          * The user wants to run a command with a set and possibly cpumask.
 291          */
 292         if (argc) {
 293                 if (Cflag || pflag || tflag || xflag || jflag)
 294                         usage();
 295                 if (sflag) {
 296                         if (cpuset_setid(CPU_WHICH_PID, -1, setid))
 297                                 err(argc, "setid");
 298                 } else {
 299                         if (cpuset(&setid))
 300                                 err(argc, "newid");
 301                 }
 302                 if (lflag) {
 303                         if (cpuset_setaffinity(level, CPU_WHICH_PID,
 304                             -1, sizeof(mask), &mask) != 0)
 305                                 err(EXIT_FAILURE, "setaffinity");
 306                 }
 307                 errno = 0;
 308                 execvp(*argv, argv);
 309                 err(errno == ENOENT ? 127 : 126, "%s", *argv);
 310         }
 311         /*
 312          * We're modifying something that presently exists.
 313          */
 314         if (Cflag && (jflag || !pflag || sflag || tflag || xflag))
 315                 usage();
 316         if (!lflag && cflag)
 317                 usage();
 318         if (!lflag && !(Cflag || sflag))
 319                 usage();
 320         /* You can only set a mask on a thread. */
 321         if (tflag && (sflag | pflag | xflag | jflag))
 322                 usage();
 323         /* You can only set a mask on an irq. */
 324         if (xflag && (jflag | pflag | sflag | tflag))
 325                 usage();
 326         if (Cflag) {
 327                 /*
 328                  * Create a new cpuset and move the specified process
 329                  * into the set.
 330                  */
 331                 if (cpuset(&setid) < 0)
 332                         err(EXIT_FAILURE, "newid");
 333                 sflag = 1;
 334         }
 335         if (pflag && sflag) {
 336                 if (cpuset_setid(CPU_WHICH_PID, pid, setid))
 337                         err(EXIT_FAILURE, "setid");
 338                 /*
 339                  * If the user specifies a set and a list we want the mask
 340                  * to effect the pid and not the set.
 341                  */
 342                 which = CPU_WHICH_PID;
 343                 id = pid;
 344         }
 345         if (lflag) {
 346                 if (cpuset_setaffinity(level, which, id, sizeof(mask),
 347                     &mask) != 0)
 348                         err(EXIT_FAILURE, "setaffinity");
 349         }
 350
 351         exit(EXIT_SUCCESS);
 352 }
 353
 354 static void
 355 usage(void)
 356 {
 357
 358         fprintf(stderr,
 359             "usage: cpuset [-l cpu-list] [-s setid] cmd ...\n");
 360         fprintf(stderr,
 361             "       cpuset [-l cpu-list] [-s setid] -p pid\n");
 362         fprintf(stderr,
 363             "       cpuset [-c] [-l cpu-list] -C -p pid\n");
 364         fprintf(stderr,
 365             "       cpuset [-c] [-l cpu-list] [-j jailid | -p pid | -t tid | -s setid | -x irq]\n");
 366         fprintf(stderr,
 367             "       cpuset -g [-cir] [-d domain | -j jailid | -p pid | -t tid | -s setid |\n"
 368             "              -x irq]\n");
 369         exit(1);
 370 }