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/numa.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(vm_domain_policy_type_t vt)
  62 {
  63
  64         switch (vt) {
  65         case VM_POLICY_NONE:
  66                 return ("none");
  67         case VM_POLICY_ROUND_ROBIN:
  68                 return ("rr");
  69         case VM_POLICY_FIXED_DOMAIN:
  70                 return ("fixed-domain");
  71         case VM_POLICY_FIXED_DOMAIN_ROUND_ROBIN:
  72                 return ("fixed-domain-rr");
  73         case VM_POLICY_FIRST_TOUCH:
  74                 return ("first-touch");
  75         case VM_POLICY_FIRST_TOUCH_ROUND_ROBIN:
  76                 return ("first-touch-rr");
  77         default:
  78                 return ("unknown");
  79         }
  80 }
  81
  82 static int
  83 parse_policy(struct vm_domain_policy_entry *vd, const char *str)
  84 {
  85
  86         if (strcmp(str, "rr") == 0) {
  87                 vd->policy = VM_POLICY_ROUND_ROBIN;
  88                 vd->domain = -1;
  89                 return (0);
  90         }
  91
  92         if (strcmp(str, "first-touch-rr") == 0) {
  93                 vd->policy = VM_POLICY_FIRST_TOUCH_ROUND_ROBIN;
  94                 vd->domain = -1;
  95                 return (0);
  96         }
  97
  98         if (strcmp(str, "first-touch") == 0) {
  99                 vd->policy = VM_POLICY_FIRST_TOUCH;
 100                 vd->domain = -1;
 101                 return (0);
 102         }
 103
 104         if (strcmp(str, "fixed-domain") == 0) {
 105                 vd->policy = VM_POLICY_FIXED_DOMAIN;
 106                 vd->domain = 0;
 107                 return (0);
 108         }
 109
 110         if (strcmp(str, "fixed-domain-rr") == 0) {
 111                 vd->policy = VM_POLICY_FIXED_DOMAIN_ROUND_ROBIN;
 112                 vd->domain = 0;
 113                 return (0);
 114         }
 115
 116         return (-1);
 117 }
 118
 119 static void
 120 usage(void)
 121 {
 122
 123         printf("usage: numactl --get [--tid/-t <tid>] [--pid/-p <pid>]\n");
 124         printf("       numactl --set [--tid=<tid>] [--pid/-p<pid>]\n");
 125         printf("                     [--mempolicy/-l <policy>] [--memdomain/"
 126             "-m <domain>]\n");
 127         printf("                     [--cpudomain/-c <domain>]\n");
 128         printf("       numactl [--mempolicy/-l <policy>] [--memdomain/-m "
 129             "<domain>]\n");
 130         printf("               [--cpudomain/-c <domain>] <cmd> ...\n");
 131
 132         exit(EX_USAGE);
 133 }
 134
 135 static int
 136 set_numa_domain_cpuaffinity(int cpu_domain, cpuwhich_t which, id_t id)
 137 {
 138         cpuset_t set;
 139         int error;
 140
 141         error = cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_DOMAIN,
 142             cpu_domain, sizeof(set), &set);
 143         if (error != 0)
 144                 err(1, "cpuset_getaffinity");
 145         error = cpuset_setaffinity(CPU_LEVEL_WHICH, which, id, sizeof(set),
 146             &set);
 147         if (error != 0)
 148                 err(1, "cpuset_setaffinity");
 149
 150         return (0);
 151 }
 152
 153 int
 154 main(int argc, char *argv[])
 155 {
 156         struct vm_domain_policy_entry vd;
 157         lwpid_t tid;
 158         pid_t pid;
 159         cpuwhich_t which;
 160         id_t id;
 161         int error;
 162         int is_set, is_get;
 163         int mem_policy_set;
 164         int ch;
 165         int cpu_domain;
 166
 167         id = -1;
 168         which = -1;
 169         is_set = 0;
 170         is_get = 0;
 171         mem_policy_set = 0;
 172         tid = -1;
 173         pid = -1;
 174         cpu_domain = -1;
 175
 176         while ((ch = getopt_long(argc, argv, "c:gl:m:p:st:", longopts,
 177             NULL)) != -1) {
 178                 switch (ch) {
 179                 case 'c':
 180                         cpu_domain = atoi(optarg);
 181                         break;
 182                 case 'g':
 183                         is_get = 1;
 184                         break;
 185                 case 'l':
 186                         if (parse_policy(&vd, optarg) != 0) {
 187                                 fprintf(stderr,
 188                                     "Could not parse policy: '%s'\n", optarg);
 189                                 exit(1);
 190                         }
 191                         mem_policy_set = 1;
 192                         break;
 193                 case 'm':
 194                         if (mem_policy_set == 0) {
 195                                 fprintf(stderr,
 196                                     "Error: set policy first before domain\n");
 197                                 exit(1);
 198                         }
 199                         vd.domain = atoi(optarg);
 200                         break;
 201                 case 'p':
 202                         pid = atoi(optarg);
 203                         break;
 204                 case 's':
 205                         is_set = 1;
 206                         break;
 207                 case 't':
 208                         tid = atoi(optarg);
 209                         break;
 210                 default:
 211                         usage();
 212                 }
 213         }
 214         argc -= optind;
 215         argv += optind;
 216
 217         /* Handle the user wishing to run a command */
 218         if (argc) {
 219                 /* Ensure that a policy was set */
 220                 if (mem_policy_set == 0) {
 221                         fprintf(stderr, "Error: no policy given\n");
 222                         usage();
 223                 }
 224
 225                 /* Set current memory process policy, will be inherited */
 226                 if (numa_setaffinity(CPU_WHICH_PID, -1, &vd) != 0)
 227                         err(1, "numa_setaffinity");
 228
 229                 /* If a CPU domain policy was given, include that too */
 230                 if (cpu_domain != -1)
 231                         (void) set_numa_domain_cpuaffinity(cpu_domain,
 232                             CPU_WHICH_PID, -1);
 233
 234                 errno = 0;
 235                 execvp(*argv, argv);
 236                 err(errno == ENOENT ? 127 : 126, "%s", *argv);
 237         }
 238
 239         /* Figure out which */
 240         if (tid != -1) {
 241                 which = CPU_WHICH_TID;
 242                 id = tid;
 243         } else if (pid != -1) {
 244                 which = CPU_WHICH_PID;
 245                 id = pid;
 246         } else {
 247                 fprintf(stderr, "Error: one of tid or pid must be given\n");
 248                 usage();
 249         }
 250
 251         /* Sanity checks */
 252         if (is_set && is_get) {
 253                 fprintf(stderr, "Error: can't set both 'set' and 'get'\n");
 254                 usage();
 255         }
 256
 257         if (is_set && ! mem_policy_set) {
 258                 fprintf(stderr, "Error: --set given, but no policy\n");
 259                 usage();
 260         }
 261
 262         /* If it's get, then get the policy and return */
 263         if (is_get) {
 264                 error = numa_getaffinity(which, id, &vd);
 265                 if (error != 0)
 266                         err(1, "numa_getaffinity");
 267                 printf("  Policy: %s; domain: %d\n",
 268                     policy_to_str(vd.policy),
 269                     vd.domain);
 270                 exit(0);
 271         }
 272
 273         /* Assume it's set */
 274
 275         /* Syscall */
 276         error = numa_setaffinity(which, id, &vd);
 277         if (error != 0)
 278                 err(1, "numa_setaffinity");
 279
 280         /* If a CPU domain policy was given, include that too */
 281         if (cpu_domain != -1)
 282                 (void) set_numa_domain_cpuaffinity(cpu_domain, which, id);
 283
 284         exit(0);
 285 }