schedutils/taskset.c

   1 /*
   2  * taskset.c - set or retrieve a task's CPU affinity
   3  *
   4  * This program is free software; you can redistribute it and/or modify
   5  * it under the terms of the GNU General Public License, version 2, as
   6  * published by the Free Software Foundation
   7  *
   8  * This program is distributed in the hope that it will be useful,
   9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11  * GNU General Public License for more details.
  12  *
  13  * You should have received a copy of the GNU General Public License along
  14  * with this program; if not, write to the Free Software Foundation, Inc.,
  15  * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  16  *
  17  * Copyright (C) 2004 Robert Love
  18  * Copyright (C) 2010 Karel Zak <kzak@redhat.com>
  19  */
  20
  21 #include <stdio.h>
  22 #include <stdlib.h>
  23 #include <unistd.h>
  24 #include <getopt.h>
  25 #include <errno.h>
  26 #include <sched.h>
  27 #include <stddef.h>
  28 #include <string.h>
  29
  30 #include "cpuset.h"
  31 #include "nls.h"
  32 #include "strutils.h"
  33 #include "xalloc.h"
  34 #include "procutils.h"
  35 #include "c.h"
  36 #include "closestream.h"
  37
  38 struct taskset {
  39         pid_t           pid;            /* task PID */
  40         cpu_set_t       *set;           /* task CPU mask */
  41         size_t          setsize;
  42         char            *buf;           /* buffer for conversion from mask to string */
  43         size_t          buflen;
  44         unsigned int    use_list:1,     /* use list rather than masks */
  45                         get_only:1;     /* print the mask, but not modify */
  46 };
  47
  48 static void __attribute__((__noreturn__)) usage(void)
  49 {
  50         FILE *out = stdout;
  51         fprintf(out,
  52                 _("Usage: %s [options] [mask | cpu-list] [pid|cmd [args...]]\n\n"),
  53                 program_invocation_short_name);
  54
  55         fputs(USAGE_SEPARATOR, out);
  56         fputs(_("Show or change the CPU affinity of a process.\n"), out);
  57         fputs(USAGE_SEPARATOR, out);
  58
  59         fprintf(out, _(
  60                 "Options:\n"
  61                 " -a, --all-tasks         operate on all the tasks (threads) for a given pid\n"
  62                 " -p, --pid               operate on existing given pid\n"
  63                 " -c, --cpu-list          display and specify cpus in list format\n"
  64                 ));
  65         printf(USAGE_HELP_OPTIONS(25));
  66
  67         fputs(USAGE_SEPARATOR, out);
  68         fprintf(out, _(
  69                 "The default behavior is to run a new command:\n"
  70                 "    %1$s 03 sshd -b 1024\n"
  71                 "You can retrieve the mask of an existing task:\n"
  72                 "    %1$s -p 700\n"
  73                 "Or set it:\n"
  74                 "    %1$s -p 03 700\n"
  75                 "List format uses a comma-separated list instead of a mask:\n"
  76                 "    %1$s -pc 0,3,7-11 700\n"
  77                 "Ranges in list format can take a stride argument:\n"
  78                 "    e.g. 0-31:2 is equivalent to mask 0x55555555\n"),
  79                 program_invocation_short_name);
  80
  81         printf(USAGE_MAN_TAIL("taskset(1)"));
  82         exit(EXIT_SUCCESS);
  83 }
  84
  85 static void print_affinity(struct taskset *ts, int isnew)
  86 {
  87         char *str, *msg;
  88
  89         if (ts->use_list) {
  90                 str = cpulist_create(ts->buf, ts->buflen, ts->set, ts->setsize);
  91                 msg = isnew ? _("pid %d's new affinity list: %s\n") :
  92                               _("pid %d's current affinity list: %s\n");
  93         } else {
  94                 str = cpumask_create(ts->buf, ts->buflen, ts->set, ts->setsize);
  95                 msg = isnew ? _("pid %d's new affinity mask: %s\n") :
  96                               _("pid %d's current affinity mask: %s\n");
  97         }
  98
  99         if (!str)
 100                 errx(EXIT_FAILURE, _("internal error: conversion from cpuset to string failed"));
 101
 102         printf(msg, ts->pid ? ts->pid : getpid(), str);
 103 }
 104
 105 static void __attribute__((__noreturn__)) err_affinity(pid_t pid, int set)
 106 {
 107         char *msg;
 108
 109         msg = set ? _("failed to set pid %d's affinity") :
 110                     _("failed to get pid %d's affinity");
 111
 112         err(EXIT_FAILURE, msg, pid ? pid : getpid());
 113 }
 114
 115 static void do_taskset(struct taskset *ts, size_t setsize, cpu_set_t *set)
 116 {
 117         /* read the current mask */
 118         if (ts->pid) {
 119                 if (sched_getaffinity(ts->pid, ts->setsize, ts->set) < 0)
 120                         err_affinity(ts->pid, 1);
 121                 print_affinity(ts, FALSE);
 122         }
 123
 124         if (ts->get_only)
 125                 return;
 126
 127         /* set new mask */
 128         if (sched_setaffinity(ts->pid, setsize, set) < 0)
 129                 err_affinity(ts->pid, 1);
 130
 131         /* re-read the current mask */
 132         if (ts->pid) {
 133                 if (sched_getaffinity(ts->pid, ts->setsize, ts->set) < 0)
 134                         err_affinity(ts->pid, 0);
 135                 print_affinity(ts, TRUE);
 136         }
 137 }
 138
 139 int main(int argc, char **argv)
 140 {
 141         cpu_set_t *new_set;
 142         pid_t pid = 0;
 143         int c, all_tasks = 0;
 144         int ncpus;
 145         size_t new_setsize, nbits;
 146         struct taskset ts;
 147
 148         static const struct option longopts[] = {
 149                 { "all-tasks",  0, NULL, 'a' },
 150                 { "pid",        0, NULL, 'p' },
 151                 { "cpu-list",   0, NULL, 'c' },
 152                 { "help",       0, NULL, 'h' },
 153                 { "version",    0, NULL, 'V' },
 154                 { NULL,         0, NULL,  0  }
 155         };
 156
 157         setlocale(LC_ALL, "");
 158         bindtextdomain(PACKAGE, LOCALEDIR);
 159         textdomain(PACKAGE);
 160         close_stdout_atexit();
 161
 162         memset(&ts, 0, sizeof(ts));
 163
 164         while ((c = getopt_long(argc, argv, "+apchV", longopts, NULL)) != -1) {
 165                 switch (c) {
 166                 case 'a':
 167                         all_tasks = 1;
 168                         break;
 169                 case 'p':
 170                         pid = strtos32_or_err(argv[argc - 1],
 171                                             _("invalid PID argument"));
 172                         break;
 173                 case 'c':
 174                         ts.use_list = 1;
 175                         break;
 176
 177                 case 'V':
 178                         print_version(EXIT_SUCCESS);
 179                 case 'h':
 180                         usage();
 181                 default:
 182                         errtryhelp(EXIT_FAILURE);
 183                 }
 184         }
 185
 186         if ((!pid && argc - optind < 2)
 187             || (pid && (argc - optind < 1 || argc - optind > 2))) {
 188                 warnx(_("bad usage"));
 189                 errtryhelp(EXIT_FAILURE);
 190         }
 191
 192         ncpus = get_max_number_of_cpus();
 193         if (ncpus <= 0)
 194                 errx(EXIT_FAILURE, _("cannot determine NR_CPUS; aborting"));
 195
 196         /*
 197          * the ts->set is always used for the sched_getaffinity call
 198          * On the sched_getaffinity the kernel demands a user mask of
 199          * at least the size of its own cpumask_t.
 200          */
 201         ts.set = cpuset_alloc(ncpus, &ts.setsize, &nbits);
 202         if (!ts.set)
 203                 err(EXIT_FAILURE, _("cpuset_alloc failed"));
 204
 205         /* buffer for conversion from mask to string */
 206         ts.buflen = 7 * nbits;
 207         ts.buf = xmalloc(ts.buflen);
 208
 209         /*
 210          * new_set is always used for the sched_setaffinity call
 211          * On the sched_setaffinity the kernel will zero-fill its
 212          * cpumask_t if the user's mask is shorter.
 213          */
 214         new_set = cpuset_alloc(ncpus, &new_setsize, NULL);
 215         if (!new_set)
 216                 err(EXIT_FAILURE, _("cpuset_alloc failed"));
 217
 218         if (argc - optind == 1)
 219                 ts.get_only = 1;
 220
 221         else if (ts.use_list) {
 222                 if (cpulist_parse(argv[optind], new_set, new_setsize, 0))
 223                         errx(EXIT_FAILURE, _("failed to parse CPU list: %s"),
 224                              argv[optind]);
 225         } else if (cpumask_parse(argv[optind], new_set, new_setsize)) {
 226                 errx(EXIT_FAILURE, _("failed to parse CPU mask: %s"),
 227                      argv[optind]);
 228         }
 229
 230         if (all_tasks && pid) {
 231                 struct proc_tasks *tasks = proc_open_tasks(pid);
 232                 while (!proc_next_tid(tasks, &ts.pid))
 233                         do_taskset(&ts, new_setsize, new_set);
 234                 proc_close_tasks(tasks);
 235         } else {
 236                 ts.pid = pid;
 237                 do_taskset(&ts, new_setsize, new_set);
 238         }
 239
 240         free(ts.buf);
 241         cpuset_free(ts.set);
 242         cpuset_free(new_set);
 243
 244         if (!pid) {
 245                 argv += optind + 1;
 246                 execvp(argv[0], argv);
 247                 errexec(argv[0]);
 248         }
 249
 250         return EXIT_SUCCESS;
 251 }