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 } 133