059320b8c8
On Linux _NPROCESSORS_CONF reports CPU threads disabled by the kernel, while it does not on FreeBSD. Flip _NPROCESSORS_ONLN to _NPROCESSORS_CONF. While it keeps reporting the same value, it will automagically unbreak should someone change the above.
133 lines
2.5 KiB
C
133 lines
2.5 KiB
C
/*-
|
|
* 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();
|
|
__builtin_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);
|
|
}
|