nghttpx: Allow units (k, m, and g) in --{read,write}-{rate,burst}

So that you can specify --read-rate=1M --read-burst=4M

src/shrpx-unittest.cc

@@ -129,6 +129,10 @@ int main(int argc, char *argv[]) {
       !CU_add_test(pSuite, "util_select_h2", shrpx::test_util_select_h2) ||
       !CU_add_test(pSuite, "util_ipv6_numeric_addr",
                    shrpx::test_util_ipv6_numeric_addr) ||
+      !CU_add_test(pSuite, "util_utos_with_unit",
+                   shrpx::test_util_utos_with_unit) ||
+      !CU_add_test(pSuite, "util_parse_uint_with_unit",
+                   shrpx::test_util_parse_uint_with_unit) ||
       !CU_add_test(pSuite, "gzip_inflate", test_nghttp2_gzip_inflate) ||
       !CU_add_test(pSuite, "ringbuf_write", nghttp2::test_ringbuf_write) ||
       !CU_add_test(pSuite, "ringbuf_iovec", nghttp2::test_ringbuf_iovec) ||

src/shrpx.cc

@@ -816,7 +816,7 @@ Performance:
   -n, --workers=<CORES>
                      Set the number of worker threads.
                      Default: )" << get_config()->num_worker << R"(
-  --read-rate=<RATE>
+  --read-rate=<SIZE>
                      Set  maximum   average  read  rate   on  frontend
                      connection.  Setting 0 to  this option means read
                      rate is unlimited.
@@ -826,7 +826,7 @@ Performance:
                      connection.  Setting 0 to  this option means read
                      burst size is unlimited.
                      Default: )" << get_config()->read_burst << R"(
-  --write-rate=<RATE>
+  --write-rate=<SIZE>
                      Set  maximum  average   write  rate  on  frontend
                      connection.  Setting 0 to this option means write
                      rate is unlimited.
@@ -836,7 +836,7 @@ Performance:
                      connection.  Setting 0 to this option means write
                      burst size is unlimited.
                      Default: )" << get_config()->write_burst << R"(
-  --worker-read-rate=<RATE>
+  --worker-read-rate=<SIZE>
                      Set  maximum   average  read  rate   on  frontend
                      connection per worker.  Setting  0 to this option
                      means  read rate  is unlimited.   Not implemented
@@ -848,7 +848,7 @@ Performance:
                      means   read  burst   size  is   unlimited.   Not
                      implemented yet.
                      Default: )" << get_config()->worker_read_burst << R"(
-  --worker-write-rate=<RATE>
+  --worker-write-rate=<SIZE>
                      Set  maximum  average   write  rate  on  frontend
                      connection per worker.  Setting  0 to this option
                      means write  rate is unlimited.   Not implemented

src/shrpx_config.cc

@@ -356,6 +356,21 @@ int parse_uint(T *dest, const char *opt, const char *optarg) {
   return 0;
 }
 
+namespace {
+template <typename T>
+int parse_uint_with_unit(T *dest, const char *opt, const char *optarg) {
+  auto n = util::parse_uint_with_unit(optarg);
+  if (n == -1) {
+    LOG(ERROR) << opt << ": bad value: '" << optarg << "'";
+    return -1;
+  }
+
+  *dest = n;
+
+  return 0;
+}
+} // namespace
+
 template <typename T>
 int parse_int(T *dest, const char *opt, const char *optarg) {
   char *end = nullptr;
@@ -879,53 +894,39 @@ int parse_config(const char *opt, const char *optarg) {
   }
 
   if (util::strieq(opt, SHRPX_OPT_READ_RATE)) {
-    return parse_uint(&mod_config()->read_rate, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->read_rate, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_READ_BURST)) {
-    int n;
+    return parse_uint_with_unit(&mod_config()->read_burst, opt, optarg);
-
-    if (parse_uint(&n, opt, optarg) != 0) {
-      return -1;
-    }
-
-    if (n == 0) {
-      LOG(ERROR) << opt << ": specify integer strictly larger than 0";
-
-      return -1;
-    }
-
-    mod_config()->read_burst = n;
-
-    return 0;
   }
 
   if (util::strieq(opt, SHRPX_OPT_WRITE_RATE)) {
-    return parse_uint(&mod_config()->write_rate, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->write_rate, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_WRITE_BURST)) {
-    return parse_uint(&mod_config()->write_burst, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->write_burst, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_WORKER_READ_RATE)) {
     LOG(WARN) << opt << ": not implemented yet";
-    return parse_uint(&mod_config()->worker_read_rate, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->worker_read_rate, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_WORKER_READ_BURST)) {
     LOG(WARN) << opt << ": not implemented yet";
-    return parse_uint(&mod_config()->worker_read_burst, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->worker_read_burst, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_WORKER_WRITE_RATE)) {
     LOG(WARN) << opt << ": not implemented yet";
-    return parse_uint(&mod_config()->worker_write_rate, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->worker_write_rate, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_WORKER_WRITE_BURST)) {
     LOG(WARN) << opt << ": not implemented yet";
-    return parse_uint(&mod_config()->worker_write_burst, opt, optarg);
+    return parse_uint_with_unit(&mod_config()->worker_write_burst, opt, optarg);
   }
 
   if (util::strieq(opt, SHRPX_OPT_NPN_LIST)) {

src/util.cc

@@ -899,6 +899,57 @@ bool ipv6_numeric_addr(const char *host) {
   return inet_pton(AF_INET6, host, dst) == 1;
 }
 
+int64_t parse_uint_with_unit(const char *s) {
+  int64_t n = 0;
+  size_t i;
+  auto len = strlen(s);
+  if (len == 0) {
+    return -1;
+  }
+  constexpr int64_t max = std::numeric_limits<int64_t>::max();
+  for (i = 0; i < len; ++i) {
+    if ('0' <= s[i] && s[i] <= '9') {
+      if (n > max / 10) {
+        return -1;
+      }
+      n *= 10;
+      if (n > max - (s[i] - '0')) {
+        return -1;
+      }
+      n += s[i] - '0';
+      continue;
+    }
+    break;
+  }
+  if (i == len) {
+    return n;
+  }
+  if (i == 0 || i + 1 != len) {
+    return -1;
+  }
+  int mul = 1;
+  switch (s[i]) {
+  case 'K':
+  case 'k':
+    mul = 1 << 10;
+    break;
+  case 'M':
+  case 'm':
+    mul = 1 << 20;
+    break;
+  case 'G':
+  case 'g':
+    mul = 1 << 30;
+    break;
+  default:
+    return -1;
+  }
+  if (n > max / mul) {
+    return -1;
+  }
+  return n * mul;
+}
+
 } // namespace util
 
 } // namespace nghttp2

src/util.h

@@ -341,6 +341,24 @@ template <typename T> std::string utos(T n) {
   return res;
 }
 
+template <typename T> std::string utos_with_unit(T n) {
+  char u = 0;
+  if (n >= (1 << 30)) {
+    u = 'G';
+    n /= (1 << 30);
+  } else if (n >= (1 << 20)) {
+    u = 'M';
+    n /= (1 << 20);
+  } else if (n >= (1 << 10)) {
+    u = 'K';
+    n /= (1 << 10);
+  }
+  if (u == 0) {
+    return utos(n);
+  }
+  return utos(n) + u;
+}
+
 extern const char UPPER_XDIGITS[];
 
 template <typename T> std::string utox(T n) {
@@ -477,6 +495,13 @@ bool check_socket_connected(int fd);
 // Returns true if |host| is IPv6 numeric address (e.g., ::1)
 bool ipv6_numeric_addr(const char *host);
 
+// Parses NULL terminated string |s| as unsigned integer and returns
+// the parsed integer.  Additionally, if |s| ends with 'k', 'm', 'g'
+// and its upper case characters, multiply the integer by 1024, 1024 *
+// 1024 and 1024 * 1024 respectively.  If there is an error, returns
+// -1.
+int64_t parse_uint_with_unit(const char *s);
+
 } // namespace util
 
 } // namespace nghttp2

src/util_test.cc

@@ -174,4 +174,37 @@ void test_util_ipv6_numeric_addr(void) {
   CU_ASSERT(!util::ipv6_numeric_addr("localhost"));
 }
 
+void test_util_utos_with_unit(void) {
+  CU_ASSERT("0" == util::utos_with_unit(0));
+  CU_ASSERT("1023" == util::utos_with_unit(1023));
+  CU_ASSERT("1K" == util::utos_with_unit(1024));
+  CU_ASSERT("1K" == util::utos_with_unit(1025));
+  CU_ASSERT("1M" == util::utos_with_unit(1 << 20));
+  CU_ASSERT("1G" == util::utos_with_unit(1 << 30));
+  CU_ASSERT("1024G" == util::utos_with_unit(1LL << 40));
+}
+
+void test_util_parse_uint_with_unit(void) {
+  CU_ASSERT(0 == util::parse_uint_with_unit("0"));
+  CU_ASSERT(1023 == util::parse_uint_with_unit("1023"));
+  CU_ASSERT(1024 == util::parse_uint_with_unit("1k"));
+  CU_ASSERT(2048 == util::parse_uint_with_unit("2K"));
+  CU_ASSERT(1 << 20 == util::parse_uint_with_unit("1m"));
+  CU_ASSERT(1 << 21 == util::parse_uint_with_unit("2M"));
+  CU_ASSERT(1 << 30 == util::parse_uint_with_unit("1g"));
+  CU_ASSERT(1LL << 31 == util::parse_uint_with_unit("2G"));
+  CU_ASSERT(9223372036854775807LL ==
+            util::parse_uint_with_unit("9223372036854775807"));
+  // check overflow case
+  CU_ASSERT(-1 == util::parse_uint_with_unit("9223372036854775808"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit("10000000000000000000"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit("9223372036854775807G"));
+  // bad characters
+  CU_ASSERT(-1 == util::parse_uint_with_unit("1.1"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit("1a"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit("a1"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit("1T"));
+  CU_ASSERT(-1 == util::parse_uint_with_unit(""));
+}
+
 } // namespace shrpx

src/util_test.h

@@ -37,6 +37,8 @@ void test_util_utox(void);
 void test_util_http_date(void);
 void test_util_select_h2(void);
 void test_util_ipv6_numeric_addr(void);
+void test_util_utos_with_unit(void);
+void test_util_parse_uint_with_unit(void);
 
 } // namespace shrpx