usr.bin/cpuset/cpuset.c

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