usr.bin/numactl/numactl.c

   1 /*
   2  * Copyright (c) 2015 Adrian Chadd <adrian@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 AUTHOR 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 AUTHOR 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/cpuset.h>
  32 #include <sys/domainset.h>
  33 #include <sys/resource.h>
  34 #include <sys/time.h>
  35
  36 #include <ctype.h>
  37 #include <err.h>
  38 #include <errno.h>
  39 #include <getopt.h>
  40 #include <libgen.h>
  41 #include <limits.h>
  42 #include <stdint.h>
  43 #include <stdio.h>
  44 #include <stdlib.h>
  45 #include <string.h>
  46 #include <sysexits.h>
  47 #include <unistd.h>
  48
  49 static struct option longopts[] = {
  50         { "tid",        required_argument,      NULL,   't' },
  51         { "pid",        required_argument,      NULL,   'p' },
  52         { "memdomain",  required_argument,      NULL,   'm' },
  53         { "cpudomain",  required_argument,      NULL,   'c' },
  54         { "mempolicy",  required_argument,      NULL,   'l' },
  55         { "set",        no_argument,    NULL,   's' },
  56         { "get",        no_argument,    NULL,   'g' },
  57         { NULL, 0, NULL, 0 }
  58 };
  59
  60 static const char *
  61 policy_to_str(int policy)
  62 {
  63
  64         switch (policy) {
  65         case DOMAINSET_POLICY_INVALID:
  66                 return ("invalid");
  67         case DOMAINSET_POLICY_ROUNDROBIN:
  68                 return ("round-robin");
  69         case DOMAINSET_POLICY_FIRSTTOUCH:
  70                 return ("first-touch");
  71         case DOMAINSET_POLICY_PREFER:
  72                 return ("prefer");
  73         default:
  74                 return ("unknown");
  75         }
  76 }
  77
  78 static void
  79 domain_print(domainset_t *mask)
  80 {
  81         int once;
  82         int bit;
  83
  84         for (once = 0, bit = 0; bit < DOMAINSET_SETSIZE; bit++) {
  85                 if (DOMAINSET_ISSET(bit, mask)) {
  86                         if (once == 0) {
  87                                 printf("%d", bit);
  88                                 once = 1;
  89                         } else
  90                                 printf(", %d", bit);
  91                 }
  92         }
  93         printf("\n");
  94 }
  95
  96 static int
  97 parse_policy(int *policy, const char *str)
  98 {
  99
 100         if (strcmp(str, "rr") == 0) {
 101                 *policy = DOMAINSET_POLICY_ROUNDROBIN;
 102                 return (0);
 103         }
 104
 105         if (strcmp(str, "first-touch-rr") == 0) {
 106                 *policy = DOMAINSET_POLICY_FIRSTTOUCH;
 107                 return (0);
 108         }
 109
 110         if (strcmp(str, "first-touch") == 0) {
 111                 *policy = DOMAINSET_POLICY_FIRSTTOUCH;
 112                 return (0);
 113         }
 114
 115         if (strcmp(str, "fixed-domain") == 0) {
 116                 *policy = DOMAINSET_POLICY_PREFER;
 117                 return (0);
 118         }
 119
 120         if (strcmp(str, "fixed-domain-rr") == 0) {
 121                 *policy = DOMAINSET_POLICY_PREFER;
 122                 return (0);
 123         }
 124
 125         return (-1);
 126 }
 127
 128 static void
 129 usage(void)
 130 {
 131
 132         printf("usage: numactl --get [--tid/-t <tid>] [--pid/-p <pid>]\n");
 133         printf("       numactl --set [--tid=<tid>] [--pid/-p<pid>]\n");
 134         printf("                     [--mempolicy/-l <policy>] [--memdomain/"
 135             "-m <domain>]\n");
 136         printf("                     [--cpudomain/-c <domain>]\n");
 137         printf("       numactl [--mempolicy/-l <policy>] [--memdomain/-m "
 138             "<domain>]\n");
 139         printf("               [--cpudomain/-c <domain>] <cmd> ...\n");
 140
 141         exit(EX_USAGE);
 142 }
 143
 144 static int
 145 set_numa_domain_cpuaffinity(int cpu_domain, cpuwhich_t which, id_t id)
 146 {
 147         cpuset_t set;
 148         int error;
 149
 150         error = cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_DOMAIN,
 151             cpu_domain, sizeof(set), &set);
 152         if (error != 0)
 153                 err(1, "cpuset_getaffinity");
 154         error = cpuset_setaffinity(CPU_LEVEL_WHICH, which, id, sizeof(set),
 155             &set);
 156         if (error != 0)
 157                 err(1, "cpuset_setaffinity");
 158
 159         return (0);
 160 }
 161
 162 /*
 163  * Attempt to maintain compatability with old style syscalls.
 164  */
 165 static int
 166 numa_setaffinity(cpuwhich_t which, id_t id, int policy, int domain)
 167 {
 168         domainset_t mask;
 169         int p;
 170
 171         DOMAINSET_ZERO(&mask);
 172         if (policy == DOMAINSET_POLICY_PREFER)
 173                 DOMAINSET_SET(domain, &mask);
 174         else if (cpuset_getdomain(CPU_LEVEL_ROOT, CPU_WHICH_PID, -1,
 175             sizeof(mask), &mask, &p) != 0)
 176                 err(EXIT_FAILURE, "getdomain");
 177         return cpuset_setdomain(CPU_LEVEL_WHICH, which, id, sizeof(mask),
 178             &mask, policy);
 179 }
 180
 181 int
 182 main(int argc, char *argv[])
 183 {
 184         lwpid_t tid;
 185         pid_t pid;
 186         cpuwhich_t which;
 187         id_t id;
 188         int error;
 189         int is_set, is_get;
 190         int mem_policy_set;
 191         int ch;
 192         int cpu_domain;
 193         int policy;
 194         int domain;
 195
 196         id = -1;
 197         which = -1;
 198         is_set = 0;
 199         is_get = 0;
 200         mem_policy_set = 0;
 201         tid = -1;
 202         pid = -1;
 203         cpu_domain = -1;
 204         domain = -1;
 205         policy = DOMAINSET_POLICY_INVALID;
 206
 207         while ((ch = getopt_long(argc, argv, "c:gl:m:p:st:", longopts,
 208             NULL)) != -1) {
 209                 switch (ch) {
 210                 case 'c':
 211                         cpu_domain = atoi(optarg);
 212                         break;
 213                 case 'g':
 214                         is_get = 1;
 215                         break;
 216                 case 'l':
 217                         if (parse_policy(&policy, optarg) != 0) {
 218                                 fprintf(stderr,
 219                                     "Could not parse policy: '%s'\n", optarg);
 220                                 exit(1);
 221                         }
 222                         mem_policy_set = 1;
 223                         break;
 224                 case 'm':
 225                         domain = atoi(optarg);
 226                         break;
 227                 case 'p':
 228                         pid = atoi(optarg);
 229                         break;
 230                 case 's':
 231                         is_set = 1;
 232                         break;
 233                 case 't':
 234                         tid = atoi(optarg);
 235                         break;
 236                 default:
 237                         usage();
 238                 }
 239         }
 240         argc -= optind;
 241         argv += optind;
 242
 243         /* Handle the user wishing to run a command */
 244         if (argc) {
 245                 /* Ensure that a policy was set */
 246                 if (mem_policy_set == 0) {
 247                         fprintf(stderr, "Error: no policy given\n");
 248                         usage();
 249                 }
 250
 251                 /* Set current memory process policy, will be inherited */
 252                 if (numa_setaffinity(CPU_WHICH_PID, -1, policy, domain) != 0)
 253                         err(1, "numa_setaffinity");
 254
 255                 /* If a CPU domain policy was given, include that too */
 256                 if (cpu_domain != -1)
 257                         (void) set_numa_domain_cpuaffinity(cpu_domain,
 258                             CPU_WHICH_PID, -1);
 259
 260                 errno = 0;
 261                 execvp(*argv, argv);
 262                 err(errno == ENOENT ? 127 : 126, "%s", *argv);
 263         }
 264
 265         /* Figure out which */
 266         if (tid != -1) {
 267                 which = CPU_WHICH_TID;
 268                 id = tid;
 269         } else if (pid != -1) {
 270                 which = CPU_WHICH_PID;
 271                 id = pid;
 272         } else {
 273                 fprintf(stderr, "Error: one of tid or pid must be given\n");
 274                 usage();
 275         }
 276
 277         /* Sanity checks */
 278         if (is_set && is_get) {
 279                 fprintf(stderr, "Error: can't set both 'set' and 'get'\n");
 280                 usage();
 281         }
 282
 283         if (is_set && ! mem_policy_set) {
 284                 fprintf(stderr, "Error: --set given, but no policy\n");
 285                 usage();
 286         }
 287
 288         /* If it's get, then get the policy and return */
 289         if (is_get) {
 290                 domainset_t mask;
 291
 292                 error = cpuset_getdomain(CPU_LEVEL_WHICH, which, id,
 293                     sizeof(mask), &mask, &policy);
 294                 if (error != 0)
 295                         err(1, "cpuset_getdomain");
 296                 printf("  Policy: %s; domain: ",
 297                     policy_to_str(policy));
 298                 domain_print(&mask);
 299                 exit(0);
 300         }
 301
 302         /* Assume it's set */
 303
 304         /* Syscall */
 305         error = numa_setaffinity(which, id, policy, domain);
 306         if (error != 0)
 307                 err(1, "numa_setaffinity");
 308
 309         /* If a CPU domain policy was given, include that too */
 310         if (cpu_domain != -1)
 311                 (void) set_numa_domain_cpuaffinity(cpu_domain, which, id);
 312
 313         exit(0);
 314 }