Add integer parsing with unit support

src/common/parse_util.cc

@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include "parse_util.h"
+
+#include <limits>
+
+// num << bit <= MAX  ->  num <= MAX >> bit
+template <typename T, typename U>
+StatusOr<T> CheckedShiftLeft(T num, U bit) {
+  if (num <= std::numeric_limits<T>::max() >> bit) {
+    return num << bit;
+  }
+
+  return {Status::NotOK, "arithmetic overflow"};
+}
+
+StatusOr<std::uint64_t> ParseSizeAndUnit(const std::string &v) {
+  auto [num, rest] = GET_OR_RET(TryParseInt<std::uint64_t>(v.c_str(), 10));
+
+  if (*rest == 0) {
+    return num;
+  } else if (Util::EqualICase(rest, "k")) {
+    return CheckedShiftLeft(num, 10);
+  } else if (Util::EqualICase(rest, "m")) {
+    return CheckedShiftLeft(num, 20);
+  } else if (Util::EqualICase(rest, "g")) {
+    return CheckedShiftLeft(num, 30);
+  } else if (Util::EqualICase(rest, "t")) {
+    return CheckedShiftLeft(num, 40);
+  } else if (Util::EqualICase(rest, "p")) {
+    return CheckedShiftLeft(num, 50);
+  }
+
+  return {Status::NotOK, "encounter unexpected unit"};
+}

src/common/parse_util.h

@@ -26,6 +26,7 @@
 #include <tuple>
 
 #include "status.h"
+#include "string_util.h"
 
 namespace details {
 
@@ -85,7 +86,7 @@ using ParseResultAndPos = std::tuple<T, const char *>;
 // base can be in {0, 2, ..., 36}, refer to strto* in standard c for more details
 template <typename T = long long>  // NOLINT
 StatusOr<ParseResultAndPos<T>> TryParseInt(const char *v, int base = 0) {
-  char *end;
+  char *end = nullptr;
 
   errno = 0;
   auto res = details::ParseIntFunc<T>::value(v, &end, base);
@@ -140,3 +141,6 @@ StatusOr<T> ParseInt(const std::string &v, NumericRange<T> range, int base = 0)
 
   return *res;
 }
+
+// available units: K, M, G, T, P
+StatusOr<std::uint64_t> ParseSizeAndUnit(const std::string &v);

src/common/string_util.cc

@@ -26,7 +26,7 @@
 
 namespace Util {
 
-const std::string Float2String(double d) {
+std::string Float2String(double d) {
   if (std::isinf(d)) {
     return d > 0 ? "inf" : "-inf";
   }
@@ -222,22 +222,24 @@ std::string StringToHex(const std::string &input) {
 
 constexpr unsigned long long expTo1024(unsigned n) { return 1ULL << (n * 10); }
 
-void BytesToHuman(char *buf, size_t size, uint64_t n) {
+std::string BytesToHuman(uint64_t n) {
   if (n < expTo1024(1)) {
-    fmt::format_to_n(buf, size, "{}B", n);
+    return fmt::format("{}B", n);
   } else if (n < expTo1024(2)) {
-    fmt::format_to_n(buf, size, "{:.2f}K", static_cast<double>(n) / expTo1024(1));
+    return fmt::format("{:.2f}K", static_cast<double>(n) / expTo1024(1));
   } else if (n < expTo1024(3)) {
-    fmt::format_to_n(buf, size, "{:.2f}M", static_cast<double>(n) / expTo1024(2));
+    return fmt::format("{:.2f}M", static_cast<double>(n) / expTo1024(2));
   } else if (n < expTo1024(4)) {
-    fmt::format_to_n(buf, size, "{:.2f}G", static_cast<double>(n) / expTo1024(3));
+    return fmt::format("{:.2f}G", static_cast<double>(n) / expTo1024(3));
   } else if (n < expTo1024(5)) {
-    fmt::format_to_n(buf, size, "{:.2f}T", static_cast<double>(n) / expTo1024(4));
+    return fmt::format("{:.2f}T", static_cast<double>(n) / expTo1024(4));
   } else if (n < expTo1024(6)) {
-    fmt::format_to_n(buf, size, "{:.2f}P", static_cast<double>(n) / expTo1024(5));
+    return fmt::format("{:.2f}P", static_cast<double>(n) / expTo1024(5));
   } else {
-    fmt::format_to_n(buf, size, "{}B", n);
+    return fmt::format("{}B", n);
   }
+
+  return {};
 }
 
 std::vector<std::string> TokenizeRedisProtocol(const std::string &value) {

src/common/string_util.h

@@ -24,10 +24,10 @@
 
 namespace Util {
 
-const std::string Float2String(double d);
+std::string Float2String(double d);
 std::string ToLower(std::string in);
 bool EqualICase(std::string_view lhs, std::string_view rhs);
-void BytesToHuman(char *buf, size_t size, uint64_t n);
+std::string BytesToHuman(uint64_t n);
 std::string Trim(std::string in, const std::string &chars);
 std::vector<std::string> Split(const std::string &in, const std::string &delim);
 std::vector<std::string> Split2KV(const std::string &in, const std::string &delim);

src/server/server.cc

@@ -37,6 +37,7 @@
 #include "storage/compaction_checker.h"
 #include "storage/redis_db.h"
 #include "storage/scripting.h"
+#include "string_util.h"
 #include "thread_util.h"
 #include "time_util.h"
 #include "tls_util.h"
@@ -820,11 +821,10 @@ void Server::GetClientsInfo(std::string *info) {
 
 void Server::GetMemoryInfo(std::string *info) {
   std::ostringstream string_stream;
-  char used_memory_rss_human[16], used_memory_lua_human[16];
   int64_t rss = Stats::GetMemoryRSS();
-  Util::BytesToHuman(used_memory_rss_human, 16, static_cast<uint64_t>(rss));
   int memory_lua = lua_gc(lua_, LUA_GCCOUNT, 0) * 1024;
-  Util::BytesToHuman(used_memory_lua_human, 16, static_cast<uint64_t>(memory_lua));
+  std::string used_memory_rss_human = Util::BytesToHuman(rss);
+  std::string used_memory_lua_human = Util::BytesToHuman(memory_lua);
   string_stream << "# Memory\r\n";
   string_stream << "used_memory_rss:" << rss << "\r\n";
   string_stream << "used_memory_human:" << used_memory_rss_human << "\r\n";

tests/cppunit/parse_util.cc

@@ -54,3 +54,18 @@ TEST(ParseUtil, ParseInt) {
   ASSERT_EQ(*ParseInt("123", {0, 123}), 123);
   ASSERT_EQ(ParseInt("124", {0, 123}).Msg(), "out of numeric range");
 }
+
+TEST(ParseUtil, ParseSizeAndUnit) {
+  ASSERT_EQ(*ParseSizeAndUnit("123"), 123);
+  ASSERT_EQ(*ParseSizeAndUnit("123K"), 123 * 1024);
+  ASSERT_EQ(*ParseSizeAndUnit("123m"), 123 * 1024 * 1024);
+  ASSERT_EQ(*ParseSizeAndUnit("123G"), 123ull << 30);
+  ASSERT_EQ(*ParseSizeAndUnit("123t"), 123ull << 40);
+  ASSERT_FALSE(ParseSizeAndUnit("123x"));
+  ASSERT_FALSE(ParseSizeAndUnit("123 t"));
+  ASSERT_FALSE(ParseSizeAndUnit("123 "));
+  ASSERT_FALSE(ParseSizeAndUnit("t"));
+  ASSERT_TRUE(ParseSizeAndUnit("16383p"));
+  ASSERT_FALSE(ParseSizeAndUnit("16384p"));
+  ASSERT_EQ(ParseSizeAndUnit("16388p").Msg(), "arithmetic overflow");
+}