git: d9d5f2c042a5 - main - cpuset: add --count

The branch main has been updated by mjg:

URL: https://cgit.FreeBSD.org/src/commit/?id=d9d5f2c042a51a9f0dd69eb1fc349efd81ffa483

commit d9d5f2c042a51a9f0dd69eb1fc349efd81ffa483
Author:     Mateusz Guzik <mjg@FreeBSD.org>
AuthorDate: 2023-02-04 17:47:41 +0000
Commit:     Mateusz Guzik <mjg@FreeBSD.org>
CommitDate: 2023-02-04 17:50:41 +0000

    cpuset: add --count
    
    Can be used to count the number of hardware threads in the cpu set.
    
    For example:
    $ cpuset -g -p $$
    pid 2440 mask: 0, 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
    pid 2440 domain policy: first-touch mask: 0, 1
    $ cpuset -g --count -p $$
    40
    
    The intent is to replace calls to sysctl hw.ncpu and kern.smp.cpus which
    can be found in the tree, which are not adequate given existence of
    cpusets.
    
    Right now only -g -p combination is supported to reduce complexity.
    As anything else errors out, this can be expanded later as needed.
    
    Differential Revision:  https://reviews.freebsd.org/D36351
---
 usr.bin/cpuset/cpuset.1 | 16 +++++++++++++++-
 usr.bin/cpuset/cpuset.c | 51 ++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/usr.bin/cpuset/cpuset.1 b/usr.bin/cpuset/cpuset.1
index 935164394b31..1a8902e1c235 100644
--- a/usr.bin/cpuset/cpuset.1
+++ b/usr.bin/cpuset/cpuset.1
@@ -25,7 +25,7 @@
 .\"
 .\" $FreeBSD$
 .\"
-.Dd July 3, 2018
+.Dd August 25, 2022
 .Dt CPUSET 1
 .Os
 .Sh NAME
@@ -57,6 +57,10 @@
 .Fl g
 .Op Fl cir
 .Op Fl d Ar domain | Fl j Ar jail | Fl p Ar pid | Fl t Ar tid | Fl s Ar setid | Fl x Ar irq
+.Nm
+.Fl g
+.Fl -count
+.Fl p Ar pid
 .Sh DESCRIPTION
 The
 .Nm
@@ -172,6 +176,16 @@ Specifies a thread id as the target of the operation.
 .It Fl x Ar irq
 Specifies an irq as the target of the operation.
 .El
+.Pp
+The long options are as follows:
+.Bl -tag -width ".Fl -count"
+.It Fl -count
+Count the number of hardware threads included in the set. Requires
+.Fl g
+and
+.Fl p
+flags.
+.El
 .Sh EXIT STATUS
 .Ex -std
 .Sh EXAMPLES
diff --git a/usr.bin/cpuset/cpuset.c b/usr.bin/cpuset/cpuset.c
index 79c6c2b6ca79..528cbc39bbc7 100644
--- a/usr.bin/cpuset/cpuset.c
+++ b/usr.bin/cpuset/cpuset.c
@@ -44,6 +44,7 @@ __FBSDID("$FreeBSD$");
 #include <ctype.h>
 #include <err.h>
 #include <errno.h>
+#include <getopt.h>
 #include <jail.h>
 #include <limits.h>
 #include <stdio.h>
@@ -52,6 +53,9 @@ __FBSDID("$FreeBSD$");
 #include <unistd.h>
 #include <string.h>
 
+/*
+ * Short opts.
+ */
 static int Cflag;
 static int cflag;
 static int dflag;
@@ -65,12 +69,24 @@ static int rflag;
 static int sflag;
 static int tflag;
 static int xflag;
+
+/*
+ * Long-only opts.
+ */
+static int count_flag;
+
 static id_t id;
 static cpulevel_t level;
 static cpuwhich_t which;
 
+#define OPT_THREADCOUNT	(CHAR_MAX + 1)
 static void usage(void);
 
+static struct option long_opts[] = {
+	{ "count", no_argument, NULL, OPT_THREADCOUNT },
+	{ NULL,	0, NULL, 0 }
+};
+
 struct numa_policy {
 	const char 	*name;
 	int		policy;
@@ -283,6 +299,18 @@ printsetid(void)
 	    levelnames[level], setid);
 }
 
+static void
+printcpucount(void)
+{
+	cpuset_t mask;
+	CPU_ZERO(&mask);
+
+	if (cpuset_getaffinity(CPU_LEVEL_ROOT, CPU_WHICH_PID, id,
+	    sizeof(mask), &mask) != 0)
+		err(EXIT_FAILURE, "getaffinity");
+	printf("%d\n", CPU_COUNT(&mask));
+}
+
 int
 main(int argc, char *argv[])
 {
@@ -300,7 +328,8 @@ main(int argc, char *argv[])
 	level = CPU_LEVEL_WHICH;
 	which = CPU_WHICH_PID;
 	id = pid = tid = setid = -1;
-	while ((ch = getopt(argc, argv, "Ccd:gij:l:n:p:rs:t:x:")) != -1) {
+	while ((ch = getopt_long(argc, argv,
+	    "Ccd:gij:l:n:p:rs:t:x:", long_opts, NULL)) != -1) {
 		switch (ch) {
 		case 'C':
 			Cflag = 1;
@@ -359,12 +388,30 @@ main(int argc, char *argv[])
 			which = CPU_WHICH_IRQ;
 			id = atoi(optarg);
 			break;
+		case OPT_THREADCOUNT:
+			count_flag = 1;
+			break;
 		default:
 			usage();
 		}
 	}
 	argc -= optind;
 	argv += optind;
+
+	/*
+	 * count requires g and p flags and is incompatible with
+	 * everything else for simplicity.
+	 */
+	if (count_flag) {
+		if (!gflag || !pflag)
+			usage();
+		if (Cflag || cflag || dflag || iflag || jflag || lflag ||
+		    nflag || rflag || sflag || tflag || xflag)
+			usage();
+		printcpucount();
+		exit(EXIT_SUCCESS);
+	}
+
 	if (gflag) {
 		if (argc || Cflag || lflag || nflag)
 			usage();
@@ -471,5 +518,7 @@ usage(void)
 	fprintf(stderr,
     "       cpuset -g [-cir]\n"
     "              [-d domain | -j jailid | -p pid | -t tid | -s setid | -x irq]\n");
+	fprintf(stderr,
+    "       cpuset -g --count -p pid\n");
 	exit(1);
 }