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