Merge pull request #5277 from grondo/optparse-size

mergify[bot] · web-flow · commit e915f2c492be · 2023-06-19T15:34:24.000Z
support `[kKMG]` suffixes in command options that take a bytes argument
diff --git a/doc/man1/flux-filemap.rst b/doc/man1/flux-filemap.rst
@@ -68,15 +68,19 @@ OPTIONS
 
 **--small-file-threshold=N**
    Set the threshold in bytes over which a regular file is mapped through
-   the distributed content cache.  Set to 0 to always use the content cache.
-   The default is 4096 (*map* subcommand only).
+   the distributed content cache. Set to 0 to always use the content cache.
+   N may be specified as a floating point number with multiplicative suffix
+   k,K=1024, M=1024\*1024, or G=1024\*1024\*1024 up to ``INT_MAX``.
+   The default is 4K (*map* subcommand only).
 
 **--disable-mmap**
    Never map a regular file through the distributed content cache.
 
 **--chunksize=N**
    Limit the content mapped blob size to N bytes.  Set to 0 for unlimited.
-   The default is 1048576 (*map* subcommand only).
+   N may be specified as a floating point number with multiplicative suffix
+   k,K=1024, M=1024\*1024, or G=1024\*1024\*1024 up to ``INT_MAX``.
+   The default is 1M (*map* subcommand only).
 
 **--direct**
    Avoid indirection through the content cache when fetching the top level
diff --git a/doc/man1/flux-ping.rst b/doc/man1/flux-ping.rst
@@ -40,9 +40,11 @@ OPTIONS
    respectively. Default: send to “*any*”.
 
 **-p, --pad**\ *=N*
-   Include in the payload a string of length *N* bytes. The payload will be
-   echoed back in the response. This option can be used to explore the
-   effect of message size on latency. Default: no padding.
+   Include in the payload a string of length *N* bytes. *N* may be a
+   floating point number with optional multiplicative suffix k,K=1024,
+   M=1024\*1024, or G=1024\*1024\*1024. The payload will be echoed back in
+   the response. This option can be used to explore the effect of message
+   size on latency. Default: no padding.
 
 **-i, --interval**\ *=N*
    Specify the delay, in seconds, between successive requests.
diff --git a/doc/man1/flux-restore.rst b/doc/man1/flux-restore.rst
@@ -54,7 +54,10 @@ OPTIONS
    store.  Performance will vary depending on the content of the archive.
 
 **--size-limit**\ =\ *SIZE*
-   Skip restoring keys that exceed SIZE bytes (default: no limit).
+   Skip restoring keys that exceed SIZE bytes (default: no limit). SIZE may
+   be specified as a floating point number with an optional multiplicative
+   suffix k or K=1024, M=1024\*1024, or G=1024\*1024\*1024 (up to
+   ``INT_MAX``).
 
 RESOURCES
 =========
diff --git a/etc/rc1 b/etc/rc1
@@ -47,7 +47,7 @@ if test $RANK -eq 0; then
             flux module load ${backingmod} truncate
         fi
         echo "restoring content from ${dumpfile}"
-        flux restore --quiet --checkpoint --size-limit=104857600 ${dumpfile}
+        flux restore --quiet --checkpoint --size-limit=100M ${dumpfile}
         if test -n "${dumplink}"; then
             rm -f ${dumplink}
         fi
diff --git a/src/cmd/builtin/filemap.c b/src/cmd/builtin/filemap.c
@@ -29,8 +29,8 @@
 #include "src/common/libfilemap/filemap.h"
 #include "src/common/libutil/fileref.h"
 
-static const int default_chunksize = 1048576;
-static const int default_small_file_threshold = 4096;
+static const char *default_chunksize = "1M";
+static const char *default_small_file_threshold = "4K";
 
 static json_t *get_list_option (optparse_t *p,
                                 const char *name,
@@ -175,10 +175,12 @@ static int subcmd_map (optparse_t *p, int ac, char *av[])
     }
     ctx.p = p;
     ctx.verbose = optparse_get_int (p, "verbose", 0);
-    ctx.chunksize = optparse_get_int (p, "chunksize", default_chunksize);
-    ctx.threshold = optparse_get_int (p,
-                                      "small-file-threshold",
-                                      default_small_file_threshold);
+    ctx.chunksize = optparse_get_size_int (p,
+                                           "chunksize",
+                                           default_chunksize);
+    ctx.threshold = optparse_get_size_int (p,
+                                           "small-file-threshold",
+                                           default_small_file_threshold);
     ctx.disable_mmap = optparse_hasopt (p, "disable-mmap");
     ctx.tags = get_list_option (p, "tags", "main");
     ctx.h = builtin_get_flux_handle (p);
@@ -375,12 +377,12 @@ static struct optparse_option map_opts[] = {
       .usage = "Change to DIR before mapping", },
     { .name = "verbose", .key = 'v', .has_arg = 2, .arginfo = "[LEVEL]",
       .usage = "Increase output detail.", },
-    { .name = "chunksize", .has_arg = 1, .arginfo = "N",
+    { .name = "chunksize", .has_arg = 1, .arginfo = "N[KMG]",
       .usage = "Limit blob size to N bytes with 0=unlimited"
-               " (default 1048576)", },
-    { .name = "small-file-threshold", .has_arg = 1, .arginfo = "N",
+               " (default 1M)", },
+    { .name = "small-file-threshold", .has_arg = 1, .arginfo = "N[KMG]",
       .usage = "Adjust the maximum size of a \"small file\" in bytes"
-               " (default 4096)", },
+               " (default 4K)", },
     { .name = "disable-mmap", .has_arg = 0,
       .usage = "Never mmap(2) files into the content cache", },
     { .name = "tags", .key = 'T', .has_arg = 1, .arginfo = "NAME,...",
diff --git a/src/cmd/builtin/restore.c b/src/cmd/builtin/restore.c
@@ -366,7 +366,7 @@ static int cmd_restore (optparse_t *p, int ac, char *av[])
         content_flags |= CONTENT_FLAG_CACHE_BYPASS;
         kvs_checkpoint_flags |= KVS_CHECKPOINT_FLAG_CACHE_BYPASS;
     }
-    blob_size_limit = optparse_get_int (p, "size-limit", 0);
+    blob_size_limit = optparse_get_size_int (p, "size-limit", "0");
 
     h = builtin_get_flux_handle (p);
     ar = restore_create (infile);
diff --git a/src/cmd/flux-ping.c b/src/cmd/flux-ping.c
@@ -316,9 +316,7 @@ int main (int argc, char *argv[])
     if (!(target = strdup (argv[optindex])))
         log_msg_exit ("out of memory");
 
-    pad_bytes = optparse_get_int (opts, "pad", 0);
-    if (pad_bytes < 0)
-        log_msg_exit ("pad must be >= 0");
+    pad_bytes = optparse_get_size_int (opts, "pad", "0");
 
     ctx.nodeid = FLUX_NODEID_ANY;
     if (optparse_hasopt (opts, "rank")) {
diff --git a/src/common/Makefile.am b/src/common/Makefile.am
@@ -111,6 +111,7 @@ libflux_optparse_la_LIBADD = \
 	$(builddir)/libczmqcontainers/libczmqcontainers.la \
 	$(builddir)/liblsd/liblsd.la \
 	$(builddir)/libutil/fsd.lo \
+	$(builddir)/libutil/parse_size.lo \
 	$(LIBPTHREAD)
 libflux_optparse_la_LDFLAGS = \
 	-Wl,--version-script=$(srcdir)/libflux-optparse.map \
diff --git a/src/common/liboptparse/optparse.c b/src/common/liboptparse/optparse.c
@@ -19,9 +19,11 @@
 #include <ctype.h>
 #include <stdarg.h>
 #include <argz.h>
+#include <math.h>
 
 #include "src/common/libczmqcontainers/czmq_containers.h"
 #include "src/common/libutil/fsd.h"
+#include "src/common/libutil/parse_size.h"
 #include "ccan/str/str.h"
 
 #include "optparse.h"
@@ -950,6 +952,59 @@ double optparse_get_duration (optparse_t *p, const char *name,
     return d;
 }
 
+uint64_t optparse_get_size (optparse_t *p,
+                            const char *name,
+                            const char *default_value)
+{
+    int n;
+    uint64_t result;
+    const char *s = NULL;
+
+    if (default_value == NULL)
+        default_value = "0";
+
+    if ((n = optparse_getopt (p, name, &s)) < 0) {
+        optparse_fatalmsg (p, 1,
+                           "%s: optparse error: no such argument '%s'\n",
+                           p->program_name, name);
+        return (uint64_t) -1;
+    }
+    if (n == 0)
+        s = default_value;
+    if (parse_size (s, &result) < 0) {
+        optparse_fatalmsg (p, 1,
+                          "%s: invalid argument for option '%s': %s: %s\n",
+                          p->program_name,
+                          name,
+                          s,
+                          strerror (errno));
+        return (uint64_t) -1;
+    }
+    return result;
+}
+
+int optparse_get_size_int (optparse_t *p,
+                           const char *name,
+                           const char *default_value)
+{
+    uint64_t val = optparse_get_size (p, name, default_value);
+    if (val == (uint64_t)-1)
+        return -1;
+    if (val > INT_MAX) {
+        const char *s;
+        optparse_getopt (p, name, &s);
+        optparse_fatalmsg (p,
+                           1,
+                           "%s: %s: value %s too large (must be < %s)\n",
+                           p->program_name,
+                           name,
+                           s,
+                           encode_size (INT_MAX+1UL));
+        return -1;
+    }
+    return (int)val;
+}
+
 const char *optparse_get_str (optparse_t *p, const char *name,
                               const char *default_value)
 {
diff --git a/src/common/liboptparse/optparse.h b/src/common/liboptparse/optparse.h
@@ -12,6 +12,7 @@
 #define _UTIL_OPTPARSE_H
 
 #include <stdbool.h>
+#include <stdint.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -381,6 +382,27 @@ int optparse_get_int (optparse_t *p, const char *name, int default_value);
 double optparse_get_duration (optparse_t *p, const char *name,
                               double default_value);
 
+/*
+ *   Return the option argument parsed as a size in bytes (or other unit)
+ *    with optional multiplicative suffix: k or K=1024, M=1024*1024,
+ *    G=1024*1024*1024, and so on for T, P, and E. The end result is
+ *    truncated and returned as uint64_t. If there was an error parsing
+ *    the size string, then the fatal error function is called.
+ *
+ *   Return default value if option was unused. default_value=NULL
+ *    is equivalent to default_value="0".
+ */
+uint64_t optparse_get_size (optparse_t *p, const char *name,
+                            const char *default_value);
+
+/*
+ *  Like optparse_get_size(), but return an int and ensure the provided
+ *   argument does not overflow INT_MAX (calls fatal error function if so).
+ */
+int optparse_get_size_int (optparse_t *p,
+                           const char *name,
+                           const char *default_value);
+
 /*
  *   Return the option argument as a double if 'name' was used,
  *    'default_value' if not.  If the option is unknown, or the argument
diff --git a/src/common/liboptparse/test/optparse.c b/src/common/liboptparse/test/optparse.c
@@ -370,25 +370,33 @@ void test_convenience_accessors (void)
 { .name = "dub", .key = 7, .has_arg = 1, .arginfo = "", .usage = "" },
 { .name = "ndb", .key = 8, .has_arg = 1, .arginfo = "", .usage = "" },
 { .name = "dur", .key = 9, .has_arg = 1, .arginfo = "", .usage = "" },
+{ .name = "size", .key = 10, .has_arg = 1, .arginfo = "", .usage = "" },
+{ .name = "sizeint", .key = 11, .has_arg = 1, .arginfo = "", .usage = "" },
         OPTPARSE_TABLE_END,
     };
 
     char *av[] = { "test", "--foo", "--baz=hello", "--mnf=7", "--neg=-4",
-                   "--dub=5.7", "--ndb=-3.2", "--dur=1.5m", NULL };
+                   "--dub=5.7", "--ndb=-3.2", "--dur=1.5m", "--size=4G",
+                   "--sizeint=1.25G", NULL };
     int ac = ARRAY_SIZE (av) - 1;
     int rc, optindex;
 
     optparse_t *p = optparse_create ("test");
     ok (p != NULL, "create object");
 
+    ok (optparse_set (p, OPTPARSE_LOG_FN, diag) == OPTPARSE_SUCCESS,
+        "optparse_set LOG_FN");
+
     rc = optparse_add_option_table (p, opts);
     ok (rc == OPTPARSE_SUCCESS, "register options");
 
-    ok (optparse_option_index (p) == -1, "optparse_option_index returns -1 before parse");
+    ok (optparse_option_index (p) == -1,
+        "optparse_option_index returns -1 before parse");
     optindex = optparse_parse_args (p, ac, av);
     ok (optindex == ac, "parse options, verify optindex");
 
-    ok (optparse_option_index (p) == optindex, "optparse_option_index works after parse");
+    ok (optparse_option_index (p) == optindex,
+        "optparse_option_index works after parse");
 
     /* hasopt
      */
@@ -461,6 +469,61 @@ void test_convenience_accessors (void)
     ok (optparse_get_duration (p, "dur", 42) == 90.,
             "get_duration returns duration arg when present");
 
+    /* get_size
+     */
+    dies_ok ({optparse_get_size (p, "no-exist", "0"); },
+            "get_size exits on unknown arg");
+    dies_ok ({optparse_get_size (p, "foo", "0"); },
+            "get_size exits on option with no argument");
+    dies_ok ({optparse_get_size (p, "baz", "0"); },
+            "get_size exits on option with wrong type argument (string)");
+    dies_ok ({optparse_get_size (p, "neg", "42"); },
+            "get_size exits on negative arg");
+    dies_ok ({optparse_get_size (p, "dur", "42"); },
+            "get_size exits on bad suffix");
+    dies_ok ({optparse_get_size (p, "bar", "1m"); },
+            "get_size exits on bad suffix in default");
+    lives_ok ({optparse_get_size (p, "size", "1k"); },
+            "get_size lives on known arg");
+
+    ok (optparse_get_size (p, "bar", "10M") == 10*1024*1024,
+            "get_size returns default argument when arg not present");
+    ok (optparse_get_size (p, "bar", NULL) == 0,
+            "get_size default_argument=NULL results in default=0 ");
+    ok (optparse_get_size (p, "mnf", "0") == 7,
+            "get_size returns arg when present");
+    ok (optparse_get_size (p, "size", "0") == 4*1024UL*1024*1024,
+            "get_size returns size arg when present");
+
+    /* get_size_int
+     */
+    dies_ok ({optparse_get_size_int (p, "no-exist", "0"); },
+            "get_size_int exits on unknown arg");
+    dies_ok ({optparse_get_size_int (p, "foo", "0"); },
+            "get_size_int exits on option with no argument");
+    dies_ok ({optparse_get_size_int (p, "baz", "0"); },
+            "get_size_int exits on option with wrong type argument (string)");
+    dies_ok ({optparse_get_size_int (p, "neg", "42"); },
+            "get_size_int exits on negative arg");
+    dies_ok ({optparse_get_size_int (p, "dur", "42"); },
+            "get_size_int exits on bad suffix");
+    dies_ok ({optparse_get_size_int (p, "bar", "1m"); },
+            "get_size_int exits on bad suffix in default");
+    dies_ok ({optparse_get_size_int (p, "size", "1M"); },
+            "get_size_int exits on value too large");
+    lives_ok ({optparse_get_size_int (p, "mnf", "1k"); },
+            "get_size_int lives on known arg");
+
+    ok (optparse_get_size_int (p, "bar", "10M") == 10*1024*1024,
+            "get_size_int returns default argument when arg not present");
+    ok (optparse_get_size_int (p, "bar", NULL) == 0,
+            "get_size_int default_argument=NULL results in default=0 ");
+    ok (optparse_get_size_int (p, "mnf", "0") == 7,
+            "get_size_int returns arg when present");
+    ok (optparse_get_size_int (p, "sizeint", "0") == 1.25*1024L*1024*1024,
+            "get_size_int returns size arg when present");
+
+
     /* get_str
      */
     dies_ok ({optparse_get_str (p, "no-exist", NULL); },
@@ -1361,9 +1424,9 @@ static void test_optional_args ()
 
 int main (int argc, char *argv[])
 {
-    plan (299);
+    plan (323);
 
-    test_convenience_accessors (); /* 36 tests */
+    test_convenience_accessors (); /* 60 tests */
     test_usage_output (); /* 46 tests */
     test_option_cb ();  /* 16 tests */
     test_errors (); /* 9 tests */
diff --git a/src/common/libutil/parse_size.c b/src/common/libutil/parse_size.c
@@ -15,6 +15,7 @@
 #include <stdint.h>
 #include <errno.h>
 #include <stdlib.h>
+#include <stdio.h>
 #include <string.h>
 #include <stdbool.h>
 #include <math.h>
@@ -132,4 +133,29 @@ int parse_size (const char *s, uint64_t *vp)
     return 0;
 }
 
+const char *encode_size (uint64_t size)
+{
+    /* Allocate a thread-local buffer to make this function easy to use
+     * in output formats (its intended use case).
+     *
+     * Note: The maximum printable digits for a double is 15 (DBL_DIG).
+     *  We also account for an optional decimal point, suffix, and required
+     *  space for NUL to get a buffer size of 18. (We ignore the fact that
+     *  a precision is specified in the %g format below for safety).
+     */
+    static __thread char buf[18];
+    const char* suffix[] = {"", "K", "M", "G", "T", "P", "E"};
+    int i = 0;
+    double value = size;
+    while (value >= 1024) {
+        value /= 1024;
+        i++;
+    }
+    /* Note: UINT64_MAX is 16E so there is no possibility that 'i' will
+     * overflow the suffix array.
+     */
+    (void) snprintf (buf, sizeof (buf), "%.8g%s", value, suffix[i]);
+    return buf;
+}
+
 // vi:ts=4 sw=4 expandtab
diff --git a/src/common/libutil/parse_size.h b/src/common/libutil/parse_size.h
diff --git a/src/common/libutil/test/parse_size.c b/src/common/libutil/test/parse_size.c
diff --git a/t/t0007-ping.t b/t/t0007-ping.t

Original file line number	Diff line number	Diff line change
`@@ -366,7 +366,7 @@ static int cmd_restore (optparse_t p, int ac, char av[])`
`366`	`366`	`content_flags \|= CONTENT_FLAG_CACHE_BYPASS;`
`367`	`367`	`kvs_checkpoint_flags \|= KVS_CHECKPOINT_FLAG_CACHE_BYPASS;`
`368`	`368`	`}`
`369`		`- blob_size_limit = optparse_get_int (p, "size-limit", 0);`
	`369`	`+ blob_size_limit = optparse_get_size_int (p, "size-limit", "0");`
`370`	`370`
`371`	`371`	`h = builtin_get_flux_handle (p);`
`372`	`372`	`ar = restore_create (infile);`