Merge pull request #6471 from garlick/rusage

add flux module stats --rusage=[self|children|thread] optional argument
flux-framework · Dec 4, 2024 · 9f43e6b · 9f43e6b
2 parents d61eccb + e2a50f0
commit 9f43e6b
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 28 deletions.
diff --git a/doc/man1/flux-module.rst b/doc/man1/flux-module.rst
@@ -105,10 +105,11 @@ the object may be customized on a module basis.
   Multiply the returned (int or double) value by the specified
   floating point value.
 
-.. option:: -R, --rusage
+.. option:: -R, --rusage=[self|children|thread]
 
   Return a JSON object representing an *rusage* structure
-  returned by :linux:man2:`getrusage`.
+  returned by :linux:man2:`getrusage`.  If specified, the optional argument
+  specifies the query target (default: self).
 
 .. option:: -c, --clear
 

diff --git a/src/cmd/flux-module.c b/src/cmd/flux-module.c
@@ -96,8 +96,10 @@ static struct optparse_option stats_opts[] =  {
     { .name = "type", .key = 't', .has_arg = 1, .arginfo = "int|double",
       .usage = "Convert JSON value to specified type",
     },
-    { .name = "rusage", .key = 'R', .has_arg = 0,
-      .usage = "Request rusage data instead of stats",
+    { .name = "rusage", .key = 'R', .has_arg = 2,
+      .arginfo = "[self|children|thread]",
+      .usage = "Request rusage data instead of stats (default: self)",
+      .flags = OPTPARSE_OPT_SHORTOPT_OPTIONAL_ARG,
     },
     { .name = "clear", .key = 'c', .has_arg = 0,
       .usage = "Clear stats on target rank",
@@ -672,10 +674,22 @@ int cmd_stats (optparse_t *p, int argc, char **argv)
         flux_msg_destroy (msg);
     } else if (optparse_hasopt (p, "rusage")) {
         topic = xasprintf ("%s.rusage", service);
-        if (!(f = flux_rpc (h, topic, NULL, nodeid, 0)))
-            log_err_exit ("%s", topic);
+        const char *optarg = optparse_get_str (p, "rusage", NULL);
+        if (optarg) {
+            if (!(f = flux_rpc_pack (h,
+                                     topic,
+                                     nodeid,
+                                     0,
+                                     "{s:s}",
+                                     "who", optarg)))
+                log_err_exit ("%s", topic);
+        }
+        else {
+            if (!(f = flux_rpc (h, topic, NULL, nodeid, 0)))
+                log_err_exit ("%s", topic);
+        }
         if (flux_rpc_get (f, &json_str) < 0)
-            log_err_exit ("%s", topic);
+            log_msg_exit ("%s: %s", topic, future_strerror (f, errno));
         if (!json_str)
             log_errn_exit (EPROTO, "%s", topic);
         parse_json (p, json_str);

diff --git a/src/common/libfluxutil/method.c b/src/common/libfluxutil/method.c
@@ -17,6 +17,8 @@
 #include <sys/resource.h>
 
 #include "src/common/libutil/errno_safe.h"
+#include "src/common/libutil/errprintf.h"
+#include "ccan/str/str.h"
 
 #include "method.h"
 
@@ -125,9 +127,29 @@ void method_rusage_cb (flux_t *h,
                        void *arg)
 {
     struct rusage ru;
+    const char *s = NULL;
+    int who;
+    const char *errmsg = NULL;
+    flux_error_t error;
 
-    if (flux_request_decode (msg, NULL, NULL) < 0
-        || getrusage (RUSAGE_THREAD, &ru) < 0)
+    if (flux_request_unpack (msg, NULL, "{s?s}", "who", &s) < 0
+        && flux_request_decode (msg, NULL, NULL) < 0)
+        goto error;
+    if (!s || streq (s, "self"))
+        who = RUSAGE_SELF;
+    else if (streq (s, "children"))
+        who = RUSAGE_CHILDREN;
+#ifdef RUSAGE_THREAD
+    else if (streq (s, "thread"))
+        who = RUSAGE_THREAD;
+#endif
+    else {
+        errprintf (&error, "%s is unsupported", s);
+        errmsg = error.text;
+        errno = EINVAL;
+        goto error;
+    }
+    if (getrusage (who, &ru) < 0)
         goto error;
     if (flux_respond_pack (h, msg,
             "{s:f s:f s:i s:i s:i s:i s:i s:i s:i s:i s:i s:i s:i s:i s:i s:i}",
@@ -150,7 +172,7 @@ void method_rusage_cb (flux_t *h,
         flux_log_error (h, "error responding to rusage request");
     return;
 error:
-    if (flux_respond_error (h, msg, errno, NULL) < 0)
+    if (flux_respond_error (h, msg, errno, errmsg) < 0)
         flux_log_error (h, "error responding to rusage request");
 }
 

diff --git a/t/t0003-module.t b/t/t0003-module.t
@@ -184,25 +184,21 @@ test_expect_success 'flux module stats --scale works' '
 	test "$EVENT_TX2" -eq $((${EVENT_TX}*2))
 '
 
-
 test_expect_success 'flux module stats --rusage works' '
-	flux module stats --rusage $REALMOD >rusage.stats &&
-	grep -q utime rusage.stats &&
-	grep -q stime rusage.stats &&
-	grep -q maxrss rusage.stats &&
-	grep -q ixrss rusage.stats &&
-	grep -q idrss rusage.stats &&
-	grep -q isrss rusage.stats &&
-	grep -q minflt rusage.stats &&
-	grep -q majflt rusage.stats &&
-	grep -q nswap rusage.stats &&
-	grep -q inblock rusage.stats &&
-	grep -q oublock rusage.stats &&
-	grep -q msgsnd rusage.stats &&
-	grep -q msgrcv rusage.stats &&
-	grep -q nsignals rusage.stats &&
-	grep -q nvcsw rusage.stats &&
-	grep -q nivcsw rusage.stats
+	flux module stats --rusage $REALMOD
+'
+test_expect_success 'flux module stats -Rself works' '
+	flux module stats -Rself $REALMOD
+'
+test_expect_success 'flux module stats --rusage=children works' '
+	flux module stats --rusage=children $REALMOD
+'
+# RUSAGE_THREAD is a non-portable GNU extension
+test_expect_success 'flux module stats --rusage=thread might work :-)' '
+	test_might_fail flux module stats --rusage=thread $REALMOD
+'
+test_expect_success 'flux module stats --rusage=badopt fails' '
+	test_must_fail flux module stats --rusage=badopt $REALMOD
 '
 
 test_expect_success 'flux module stats --rusage --parse maxrss works' '