usr.bin/cpuset/cpuset.c

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