fix parseDisjointPoolConfig and add tests

Author: bratpiorka

source/common/umf_pools/disjoint_pool_config_parser.cpp

@@ -72,151 +72,35 @@ DisjointPoolAllConfigs::DisjointPoolAllConfigs(int trace) {
 
 DisjointPoolAllConfigs parseDisjointPoolConfig(const std::string &config,
                                                int trace) {
+    // TODO: avoid creating a copy of 'config'
     DisjointPoolAllConfigs AllConfigs;
 
-    // TODO: replace with UR ENV var parser and avoid creating a copy of 'config'
-    auto GetValue = [](std::string &Param, size_t Length, size_t &Setting) {
-        size_t Multiplier = 1;
-        if (tolower(Param[Length - 1]) == 'k') {
-            Length--;
-            Multiplier = 1_KB;
-        }
-        if (tolower(Param[Length - 1]) == 'm') {
-            Length--;
-            Multiplier = 1_MB;
-        }
-        if (tolower(Param[Length - 1]) == 'g') {
-            Length--;
-            Multiplier = 1_GB;
-        }
-        std::string TheNumber = Param.substr(0, Length);
-        if (TheNumber.find_first_not_of("0123456789") == std::string::npos) {
-            Setting = std::stoi(TheNumber) * Multiplier;
-        }
-    };
+    size_t MaxSize = (std::numeric_limits<size_t>::max)();
+    size_t EnableBuffers = 1;
 
-    auto ParamParser = [GetValue](std::string &Params, size_t &Setting,
-                                  bool &ParamWasSet) {
-        bool More;
-        if (Params.size() == 0) {
-            ParamWasSet = false;
-            return false;
-        }
-        size_t Pos = Params.find(',');
-        if (Pos != std::string::npos) {
-            if (Pos > 0) {
-                GetValue(Params, Pos, Setting);
-                ParamWasSet = true;
-            }
-            Params.erase(0, Pos + 1);
-            More = true;
+    // Update pool settings if specified in environment.
+    bool EnableBuffersSet = false;
+    bool MaxSizeSet = false;
+    size_t Start = 0;
+    size_t End = config.find(';');
+    while (true) {
+        std::string Param = config.substr(Start, End - Start);
+        if (!EnableBuffersSet && isdigit(Param[0])) {
+            GetValue(Param, Param.size(), EnableBuffers);
+            EnableBuffersSet = true;
+        } else if (!MaxSizeSet && isdigit(Param[0])) {
+            GetValue(Param, Param.size(), MaxSize);
+            MaxSizeSet = true;
         } else {
-            GetValue(Params, Params.size(), Setting);
-            ParamWasSet = true;
-            More = false;
+            EnvVarMap map = string_to_map
         }
-        return More;
-    };
 
-    auto MemParser = [&AllConfigs, ParamParser](std::string &Params,
-                                                DisjointPoolMemType memType =
-                                                    DisjointPoolMemType::All) {
-        bool ParamWasSet;
-        DisjointPoolMemType LM = memType;
-        if (memType == DisjointPoolMemType::All) {
-            LM = DisjointPoolMemType::Host;
+        if (End == std::string::npos) {
+            break;
         }
 
-        bool More = ParamParser(Params, AllConfigs.Configs[LM].MaxPoolableSize,
-                                ParamWasSet);
-        if (ParamWasSet && memType == DisjointPoolMemType::All) {
-            for (auto &Config : AllConfigs.Configs) {
-                Config.MaxPoolableSize = AllConfigs.Configs[LM].MaxPoolableSize;
-            }
-        }
-        if (More) {
-            More = ParamParser(Params, AllConfigs.Configs[LM].Capacity,
-                               ParamWasSet);
-            if (ParamWasSet && memType == DisjointPoolMemType::All) {
-                for (auto &Config : AllConfigs.Configs) {
-                    Config.Capacity = AllConfigs.Configs[LM].Capacity;
-                }
-            }
-        }
-        if (More) {
-            ParamParser(Params, AllConfigs.Configs[LM].SlabMinSize,
-                        ParamWasSet);
-            if (ParamWasSet && memType == DisjointPoolMemType::All) {
-                for (auto &Config : AllConfigs.Configs) {
-                    Config.SlabMinSize = AllConfigs.Configs[LM].SlabMinSize;
-                }
-            }
-        }
-    };
-
-    auto MemTypeParser = [MemParser](std::string &Params) {
-        int Pos = 0;
-        DisjointPoolMemType M(DisjointPoolMemType::All);
-        if (Params.compare(0, 5, "host:") == 0) {
-            Pos = 5;
-            M = DisjointPoolMemType::Host;
-        } else if (Params.compare(0, 7, "device:") == 0) {
-            Pos = 7;
-            M = DisjointPoolMemType::Device;
-        } else if (Params.compare(0, 7, "shared:") == 0) {
-            Pos = 7;
-            M = DisjointPoolMemType::Shared;
-        } else if (Params.compare(0, 17, "read_only_shared:") == 0) {
-            Pos = 17;
-            M = DisjointPoolMemType::SharedReadOnly;
-        }
-        if (Pos > 0) {
-            Params.erase(0, Pos);
-        }
-        MemParser(Params, M);
-    };
-
-    size_t MaxSize = (std::numeric_limits<size_t>::max)();
-
-    // Update pool settings if specified in environment.
-    size_t EnableBuffers = 1;
-    if (config != "") {
-        std::string Params = config;
-        size_t Pos = Params.find(';');
-        if (Pos != std::string::npos) {
-            if (Pos > 0) {
-                GetValue(Params, Pos, EnableBuffers);
-            }
-            Params.erase(0, Pos + 1);
-            size_t Pos = Params.find(';');
-            if (Pos != std::string::npos) {
-                if (Pos > 0) {
-                    GetValue(Params, Pos, MaxSize);
-                }
-                Params.erase(0, Pos + 1);
-                do {
-                    size_t Pos = Params.find(';');
-                    if (Pos != std::string::npos) {
-                        if (Pos > 0) {
-                            std::string MemParams = Params.substr(0, Pos);
-                            MemTypeParser(MemParams);
-                        }
-                        Params.erase(0, Pos + 1);
-                        if (Params.size() == 0) {
-                            break;
-                        }
-                    } else {
-                        MemTypeParser(Params);
-                        break;
-                    }
-                } while (true);
-            } else {
-                // set MaxPoolSize for all configs
-                GetValue(Params, Params.size(), MaxSize);
-            }
-        } else {
-            GetValue(Params, Params.size(), EnableBuffers);
-        }
+        Start = End + 1;
+        End = config.find(';', Start);
     }
 
     AllConfigs.EnableBuffers = EnableBuffers;

source/common/ur_util.hpp

@@ -209,18 +209,27 @@ using EnvVarMap = std::map<std::string, std::vector<std::string>>;
 ///         Otherwise, optional is set to std::nullopt when the environment variable
 ///         is not set or is empty.
 /// @throws std::invalid_argument() when the parsed environment variable has wrong format
+
 inline std::optional<EnvVarMap> getenv_to_map(const char *env_var_name,
                                               bool reject_empty = true) {
-    char main_delim = ';';
-    char key_value_delim = ':';
-    char values_delim = ',';
-    EnvVarMap map;
-
     auto env_var = ur_getenv(env_var_name);
     if (!env_var.has_value()) {
         return std::nullopt;
     }
 
+    try {
+        auto map = string_to_map(*env_var, reject_empty);
+    } catch (...) {
+        throw_wrong_format_map(env_var_name, env_var);
+    }
+}
+
+inline EnvVarMap string_to_map(std::string input, bool reject_empty = true) {
+    char main_delim = ';';
+    char key_value_delim = ':';
+    char values_delim = ',';
+    EnvVarMap map;
+
     auto is_quoted = [](std::string &str) {
         return (str.front() == '\'' && str.back() == '\'') ||
                (str.front() == '"' && str.back() == '"');
@@ -229,30 +238,30 @@ inline std::optional<EnvVarMap> getenv_to_map(const char *env_var_name,
         return str.find(':') != std::string::npos;
     };
 
-    std::stringstream ss(*env_var);
+    std::stringstream ss(input);
     std::string key_value;
     while (std::getline(ss, key_value, main_delim)) {
         std::string key;
         std::string values;
         std::stringstream kv_ss(key_value);
 
         if (reject_empty && !has_colon(key_value)) {
-            throw_wrong_format_map(env_var_name, *env_var);
+            throw_wrong_format_map(env_var_name, input);
         }
 
         std::getline(kv_ss, key, key_value_delim);
         std::getline(kv_ss, values);
         if (key.empty() || (reject_empty && values.empty()) ||
             map.find(key) != map.end()) {
-            throw_wrong_format_map(env_var_name, *env_var);
+            throw;
         }
 
         std::vector<std::string> values_vec;
         std::stringstream values_ss(values);
         std::string value;
         while (std::getline(values_ss, value, values_delim)) {
             if (value.empty() || (has_colon(value) && !is_quoted(value))) {
-                throw_wrong_format_map(env_var_name, *env_var);
+                throw;
             }
             if (is_quoted(value)) {
                 value.erase(value.cbegin());

test/usm/usmPoolManager.cpp

@@ -3,6 +3,7 @@
 // See LICENSE.TXT
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "umf_pools/disjoint_pool_config_parser.hpp"
 #include "ur_pool_manager.hpp"
 
 #include <uur/fixtures.h>
@@ -18,6 +19,27 @@ auto createMockPoolHandle() {
         [](umf_memory_pool_t *) {});
 }
 
+bool compareConfig(const usm::umf_disjoint_pool_config_t &left,
+                   usm::umf_disjoint_pool_config_t &right) {
+    return left.MaxPoolableSize == right.MaxPoolableSize &&
+           left.Capacity == right.Capacity &&
+           left.SlabMinSize == right.SlabMinSize;
+}
+
+bool compareConfigs(const usm::DisjointPoolAllConfigs &left,
+                    usm::DisjointPoolAllConfigs &right) {
+    return left.EnableBuffers == right.EnableBuffers &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Host],
+                         right.Configs[usm::DisjointPoolMemType::Host]) &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Device],
+                         right.Configs[usm::DisjointPoolMemType::Device]) &&
+           compareConfig(left.Configs[usm::DisjointPoolMemType::Shared],
+                         right.Configs[usm::DisjointPoolMemType::Shared]) &&
+           compareConfig(
+               left.Configs[usm::DisjointPoolMemType::SharedReadOnly],
+               right.Configs[usm::DisjointPoolMemType::SharedReadOnly]);
+}
+
 TEST_P(urUsmPoolDescriptorTest, poolIsPerContextTypeAndDevice) {
     auto &devices = uur::DevicesEnvironment::instance->devices;
     auto poolHandle = this->GetParam();
@@ -115,4 +137,29 @@ TEST_P(urUsmPoolManagerTest, poolManagerGetNonexistant) {
     }
 }
 
+TEST_P(urUsmPoolManagerTest, config) {
+    // Check default config
+    usm::DisjointPoolAllConfigs def;
+    usm::DisjointPoolAllConfigs parsed1 =
+        usm::parseDisjointPoolConfig("1;host:2M,4,64K;device:4M,4,64K;"
+                                     "shared:0,0,2M;read_only_shared:4M,4,2M",
+                                     0);
+    ASSERT_EQ(compareConfigs(def, parsed1), true);
+
+    // Check partially set config
+    usm::DisjointPoolAllConfigs parsed2 =
+        usm::parseDisjointPoolConfig("1;device:4M;shared:0,0,2M", 0);
+    ASSERT_EQ(compareConfigs(def, parsed2), true);
+
+    // Check non-default config
+    usm::DisjointPoolAllConfigs test(def);
+    test.Configs[usm::DisjointPoolMemType::Shared].MaxPoolableSize = 128 * 1024;
+    test.Configs[usm::DisjointPoolMemType::Shared].Capacity = 4;
+    test.Configs[usm::DisjointPoolMemType::Shared].SlabMinSize = 64 * 1024;
+
+    usm::DisjointPoolAllConfigs parsed3 =
+        usm::parseDisjointPoolConfig("1;shared:128K,4,64K", 0);
+    ASSERT_EQ(compareConfigs(test, parsed3), true);
+}
+
 UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urUsmPoolManagerTest);