From cfc638f05168df60ba223cd1700a8340e8e0db2c Mon Sep 17 00:00:00 2001
From: Yukang-Lian <Yukang.Lian2022@gmail.com>
Date: Thu, 30 Nov 2023 15:13:46 +0800
Subject: [PATCH] 7

---
 be/src/olap/wal_manager.cpp       | 42 +++++++++----------------------
 be/test/olap/wal_manager_test.cpp | 23 +++++++++++++++++
 2 files changed, 35 insertions(+), 30 deletions(-)

diff --git a/be/src/olap/wal_manager.cpp b/be/src/olap/wal_manager.cpp
index 4ff396290ace116..bad9f803150f548 100644
--- a/be/src/olap/wal_manager.cpp
+++ b/be/src/olap/wal_manager.cpp
@@ -34,6 +34,7 @@
 #include "runtime/fragment_mgr.h"
 #include "runtime/plan_fragment_executor.h"
 #include "runtime/stream_load/stream_load_context.h"
+#include "util/parse_util.h"
 #include "util/path_util.h"
 #include "util/thrift_rpc_helper.h"
 #include "vec/exec/format/wal/wal_reader.h"
@@ -84,36 +85,17 @@ Status WalManager::init() {
 }
 
 Status WalManager::init_wal_limit() {
-    const std::string& wal_limit_config = config::wal_max_disk_limit;
-    if (*wal_limit_config.rbegin() == '%') {
-        // Config is percentage.
-        size_t available_bytes;
-        size_t disk_capacity_bytes;
-        // Get the root path available space.
-        RETURN_IF_ERROR(io::global_local_filesystem()->get_space_info("./", &disk_capacity_bytes,
-                                                                      &available_bytes));
-        StringParser::ParseResult result;
-        auto limit_val = StringParser::string_to_float<double>(
-                wal_limit_config.data(), wal_limit_config.size() - 1, &result);
-        if (result != StringParser::PARSE_SUCCESS || limit_val < 0 || limit_val > 100) {
-            return Status::InternalError(
-                    "Parse wal_max_disk_limit=xx% wrong, please check your config!");
-        }
-        wal_limit = available_bytes * (static_cast<double>(limit_val) / 100.0);
-    } else {
-        if (*wal_limit_config.begin() == '-') {
-            return Status::InternalError(
-                    "wal_max_disk_limit should be positive, please check your config!");
-        }
-        // Config is number.
-        StringParser::ParseResult result;
-        auto limit_val = StringParser::string_to_int<size_t>(wal_limit_config.data(),
-                                                             wal_limit_config.size(), &result);
-        if (result != StringParser::PARSE_SUCCESS) {
-            return Status::InternalError(
-                    "Parse wal_max_disk_limit=xx wrong, please check your config!");
-        }
-        wal_limit = limit_val;
+    size_t available_bytes;
+    size_t disk_capacity_bytes;
+    // Get the root path available space.
+    RETURN_IF_ERROR(io::global_local_filesystem()->get_space_info("./", &disk_capacity_bytes,
+                                                                  &available_bytes));
+    bool is_percent = true;
+    int64_t wal_disk_limit =
+            ParseUtil::parse_mem_spec(config::wal_max_disk_limit, -1, available_bytes, &is_percent);
+    if (wal_disk_limit < 0) {
+        return Status::InternalError(
+                "wal_max_disk_limit config is wrong, please check your config!");
     }
     return Status::OK();
 }
diff --git a/be/test/olap/wal_manager_test.cpp b/be/test/olap/wal_manager_test.cpp
index 5ba42c4587a219d..584c86e7eb25c01 100644
--- a/be/test/olap/wal_manager_test.cpp
+++ b/be/test/olap/wal_manager_test.cpp
@@ -187,9 +187,32 @@ TEST_F(WalManagerTest, TestDynamicWalSpaceLimt) {
     EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::OK());
     EXPECT_EQ(WalManager::wal_limit, 1073741824000);
 
+    // 1M
+    WalManager::wal_limit = 0;
+    config::wal_max_disk_limit = "1M";
+    EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::OK());
+    EXPECT_EQ(WalManager::wal_limit, 1048576);
+
+    // 1G
+    WalManager::wal_limit = 0;
+    config::wal_max_disk_limit = "1G";
+    EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::OK());
+    EXPECT_EQ(WalManager::wal_limit, 1073741824);
+
+    // 100G
+    WalManager::wal_limit = 0;
+    config::wal_max_disk_limit = "100G";
+    EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::OK());
+    EXPECT_EQ(WalManager::wal_limit, 1073741824000);
+
     WalManager::wal_limit = 0;
     config::wal_max_disk_limit = "-1024";
     EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::InternalError(""));
     EXPECT_EQ(WalManager::wal_limit, 0);
+
+    WalManager::wal_limit = 0;
+    config::wal_max_disk_limit = "-1M";
+    EXPECT_EQ(_env->wal_mgr()->init_wal_limit(), Status::InternalError(""));
+    EXPECT_EQ(WalManager::wal_limit, 0);
 }
 } // namespace doris
\ No newline at end of file