busybox/coreutils/nproc.c
<<
>>
Prefs
   1/*
   2 * Copyright (C) 2017 Denys Vlasenko <vda.linux@googlemail.com>
   3 *
   4 * Licensed under GPLv2, see LICENSE in this source tree
   5 */
   6//config:config NPROC
   7//config:       bool "nproc (3.7 kb)"
   8//config:       default y
   9//config:       help
  10//config:       Print number of CPUs
  11
  12//applet:IF_NPROC(APPLET_NOFORK(nproc, nproc, BB_DIR_USR_BIN, BB_SUID_DROP, nproc))
  13
  14//kbuild:lib-$(CONFIG_NPROC) += nproc.o
  15
  16//usage:#define nproc_trivial_usage
  17//usage:        ""IF_LONG_OPTS("[--all] [--ignore=N]")
  18//usage:#define nproc_full_usage "\n\n"
  19//usage:        "Print number of available CPUs"
  20//usage:        IF_LONG_OPTS(
  21//usage:     "\n"
  22//usage:     "\n        --all           Number of installed CPUs"
  23//usage:     "\n        --ignore=N      Exclude N CPUs"
  24//usage:        )
  25
  26#include <sched.h>
  27#include "libbb.h"
  28
  29int nproc_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
  30int nproc_main(int argc UNUSED_PARAM, char **argv UNUSED_PARAM)
  31{
  32        unsigned long mask[1024];
  33        int count = 0;
  34#if ENABLE_LONG_OPTS
  35        int ignore = 0;
  36        int opts = getopt32long(argv, "\xfe:+",
  37                        "ignore\0" Required_argument "\xfe"
  38                        "all\0"    No_argument       "\xff"
  39                        , &ignore
  40        );
  41
  42        if (opts & (1 << 1)) {
  43                DIR *cpusd = opendir("/sys/devices/system/cpu");
  44                if (cpusd) {
  45                        struct dirent *de;
  46                        while (NULL != (de = readdir(cpusd))) {
  47                                char *cpuid = strstr(de->d_name, "cpu");
  48                                if (cpuid && isdigit(cpuid[strlen(cpuid) - 1]))
  49                                        count++;
  50                        }
  51                        IF_FEATURE_CLEAN_UP(closedir(cpusd);)
  52                }
  53        } else
  54#endif
  55        if (sched_getaffinity(0, sizeof(mask), (void*)mask) == 0) {
  56                int i;
  57                for (i = 0; i < ARRAY_SIZE(mask); i++) {
  58                        unsigned long m = mask[i];
  59                        while (m) {
  60                                if (m & 1)
  61                                        count++;
  62                                m >>= 1;
  63                        }
  64                }
  65        }
  66
  67        IF_LONG_OPTS(count -= ignore;)
  68        if (count <= 0)
  69                count = 1;
  70
  71        printf("%u\n", count);
  72
  73        return 0;
  74}
  75