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
help(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
usage(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
version(void)68 version(void)
69 {
70 printf("nproc (neither_GNU nor_coreutils) 8.32\n");
71 exit(EXIT_SUCCESS);
72 }
73
74 int
main(int argc,char * argv[])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