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