summaryrefslogtreecommitdiff
path: root/bin/nproc/nproc.c
blob: d8affe11b79670ed337e744979d1063df6b7973c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
/*-
 * Copyright (c) 2023 Mateusz Guzik
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

/*
 * This program is intended to be compatible with nproc as found in GNU
 * coreutils.
 *
 * In order to maintain that, do not add any features here if they are not
 * present in said program.  If you are looking for anything more advanced you
 * probably should patch cpuset(1) instead.
 */

#include <sys/param.h>
#include <sys/cpuset.h>

#include <err.h>
#include <errno.h>
#include <getopt.h>
#include <limits.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <sysexits.h>
#include <unistd.h>

#define OPT_ALL		(CHAR_MAX + 1)
#define OPT_IGNORE	(CHAR_MAX + 2)
#define OPT_VERSION	(CHAR_MAX + 3)
#define OPT_HELP	(CHAR_MAX + 4)

static struct option long_opts[] = {
	{ "all", no_argument, NULL, OPT_ALL },
	{ "ignore", required_argument, NULL, OPT_IGNORE },
	{ "version", no_argument, NULL, OPT_VERSION },
	{ "help", no_argument, NULL, OPT_HELP },
	{ NULL, 0, NULL, 0 }
};

static void
help(void)
{
	fprintf(stderr,
    "usage: nproc [--all] [--ignore=count]\n");
	fprintf(stderr,
    "       nproc --help\n");
	fprintf(stderr,
    "       nproc --version\n");
}

static void
usage(void)
{
	help();
	exit(EX_USAGE);
}

/*
 * GNU variant ships with the --version switch.
 *
 * While we don't have anything to put there, print something which is
 * whitespace-compatible with the original. Version number was taken
 * from coreutils this code is in sync with.
 */
static void
version(void)
{
	printf("nproc (neither_GNU nor_coreutils) 8.32\n");
	exit(EXIT_SUCCESS);
}

int
main(int argc, char *argv[])
{
	const char *errstr;
	cpuset_t mask;
	int ch, cpus, ignore;
	bool all_flag;

	ignore = 0;
	all_flag = false;

	while ((ch = getopt_long(argc, argv, "", long_opts, NULL)) != -1) {
		switch (ch) {
		case OPT_ALL:
			all_flag = true;
			break;
		case OPT_IGNORE:
			ignore = strtonum(optarg, 0, INT_MAX, &errstr);
			if (errstr)
				errx(1, "bad ignore count: %s", errstr);
			break;
		case OPT_VERSION:
			version();
			__unreachable();
		case OPT_HELP:
			help();
			exit(EXIT_SUCCESS);
		default:
			usage();
		}
	}

	argc -= optind;
	argv += optind;

	if (argc != 0)
		usage();

	if (all_flag) {
		cpus = sysconf(_SC_NPROCESSORS_CONF);
		if (cpus == -1)
			err(1, "sysconf");
	} else {
		CPU_ZERO(&mask);
		if (cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_TID, -1,
		    sizeof(mask), &mask) != 0)
			err(1, "cpuset_getaffinity");
		cpus = CPU_COUNT(&mask);
	}

	if (ignore >= cpus)
		cpus = 1;
	else
		cpus -= ignore;

	printf("%u\n", cpus);

	exit(EXIT_SUCCESS);
}