bin/nproc/nproc.c

   1 /*-
   2  * Copyright (c) 2023 Mateusz Guzik
   3  *
   4  * SPDX-License-Identifier: BSD-2-Clause
   5  */
   6
   7 /*
   8  * This program is intended to be compatible with nproc as found in GNU
   9  * coreutils.
  10  *
  11  * In order to maintain that, do not add any features here if they are not
  12  * present in said program.  If you are looking for anything more advanced you
  13  * probably should patch cpuset(1) instead.
  14  */
  15
  16 #include <sys/param.h>
  17 #include <sys/cpuset.h>
  18
  19 #include <err.h>
  20 #include <errno.h>
  21 #include <getopt.h>
  22 #include <limits.h>
  23 #include <stdbool.h>
  24 #include <stdio.h>
  25 #include <stdlib.h>
  26 #include <sysexits.h>
  27 #include <unistd.h>
  28
  29 #define OPT_ALL         (CHAR_MAX + 1)
  30 #define OPT_IGNORE      (CHAR_MAX + 2)
  31 #define OPT_VERSION     (CHAR_MAX + 3)
  32 #define OPT_HELP        (CHAR_MAX + 4)
  33
  34 static struct option long_opts[] = {
  35         { "all", no_argument, NULL, OPT_ALL },
  36         { "ignore", required_argument, NULL, OPT_IGNORE },
  37         { "version", no_argument, NULL, OPT_VERSION },
  38         { "help", no_argument, NULL, OPT_HELP },
  39         { NULL, 0, NULL, 0 }
  40 };
  41
  42 static void
  43 help(void)
  44 {
  45         fprintf(stderr,
  46     "usage: nproc [--all] [--ignore=count]\n");
  47         fprintf(stderr,
  48     "       nproc --help\n");
  49         fprintf(stderr,
  50     "       nproc --version\n");
  51 }
  52
  53 static void
  54 usage(void)
  55 {
  56         help();
  57         exit(EX_USAGE);
  58 }
  59
  60 /*
  61  * GNU variant ships with the --version switch.
  62  *
  63  * While we don't have anything to put there, print something which is
  64  * whitespace-compatible with the original. Version number was taken
  65  * from coreutils this code is in sync with.
  66  */
  67 static void
  68 version(void)
  69 {
  70         printf("nproc (neither_GNU nor_coreutils) 8.32\n");
  71         exit(EXIT_SUCCESS);
  72 }
  73
  74 int
  75 main(int argc, char *argv[])
  76 {
  77         const char *errstr;
  78         cpuset_t mask;
  79         int ch, cpus, ignore;
  80         bool all_flag;
  81
  82         ignore = 0;
  83         all_flag = false;
  84
  85         while ((ch = getopt_long(argc, argv, "", long_opts, NULL)) != -1) {
  86                 switch (ch) {
  87                 case OPT_ALL:
  88                         all_flag = true;
  89                         break;
  90                 case OPT_IGNORE:
  91                         ignore = strtonum(optarg, 0, INT_MAX, &errstr);
  92                         if (errstr)
  93                                 errx(1, "bad ignore count: %s", errstr);
  94                         break;
  95                 case OPT_VERSION:
  96                         version();
  97                         __unreachable();
  98                 case OPT_HELP:
  99                         help();
 100                         exit(EXIT_SUCCESS);
 101                 default:
 102                         usage();
 103                 }
 104         }
 105
 106         argc -= optind;
 107         argv += optind;
 108
 109         if (argc != 0)
 110                 usage();
 111
 112         if (all_flag) {
 113                 cpus = sysconf(_SC_NPROCESSORS_CONF);
 114                 if (cpus == -1)
 115                         err(1, "sysconf");
 116         } else {
 117                 CPU_ZERO(&mask);
 118                 if (cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_TID, -1,
 119                     sizeof(mask), &mask) != 0)
 120                         err(1, "cpuset_getaffinity");
 121                 cpus = CPU_COUNT(&mask);
 122         }
 123
 124         if (ignore >= cpus)
 125                 cpus = 1;
 126         else
 127                 cpus -= ignore;
 128
 129         printf("%u\n", cpus);
 130
 131         exit(EXIT_SUCCESS);
 132 }