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