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 34static 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 42static void 43help(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 53static void 54usage(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 */ 67static void 68version(void) 69{ 70 printf("nproc (neither_GNU nor_coreutils) 8.32\n"); 71 exit(EXIT_SUCCESS); 72} 73 74int 75main(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