From f9e68ab8d06dca4a104b3e535cd65277f265bd86 Mon Sep 17 00:00:00 2001 From: Schatten Date: Mon, 31 Aug 2026 00:37:24 +0800 Subject: [PATCH] [Store] Structure LocalHotCacheConfig environment settings Signed-off-by: Schatten --- .../include/environment_variables.h | 9 ++ mooncake-store/include/client_service.h | 19 +-- mooncake-store/include/local_hot_cache.h | 9 ++ mooncake-store/src/CMakeLists.txt | 1 + mooncake-store/src/client_service.cpp | 110 ++----------- .../src/config/local_hot_cache_config.cpp | 92 +++++++++++ mooncake-store/tests/CMakeLists.txt | 1 + .../tests/local_hot_cache_config_test.cpp | 150 ++++++++++++++++++ 8 files changed, 275 insertions(+), 116 deletions(-) create mode 100644 mooncake-store/src/config/local_hot_cache_config.cpp create mode 100644 mooncake-store/tests/local_hot_cache_config_test.cpp diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 90739e9ec6..176b13714f 100644 --- a/mooncake-common/include/environment_variables.h +++ b/mooncake-common/include/environment_variables.h @@ -44,6 +44,15 @@ struct ClientAutoPortEnvironmentVariables { MC_DEFINE_ENV_VAR(int, MC_STORE_CLIENT_MAX_PORT); }; +struct LocalHotCacheEnvironmentVariables { + // Keep these values as strings to preserve their existing per-setting + // parsing, fallback, and logging behavior. + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_CACHE_SIZE); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_BLOCK_SIZE); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_CACHE_USE_SHM); + MC_DEFINE_ENV_VAR(std::string, MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD); +}; + #undef MC_DEFINE_ENV_VAR } // namespace mooncake diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index cdb156303c..4b753c0667 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -814,8 +814,8 @@ class Client { const DiskDescriptor& disk_descriptor); /** * @brief Initialize local hot cache - * @return ErrorCode::OK if use local hot cache, - * ErrorCode::INVALID_PARAMS if invalid MC_STORE_LOCAL_HOT_CACHE_SIZE config + * @return ErrorCode::OK if disabled or initialized successfully; + * ErrorCode::INVALID_PARAMS if cache allocation or registration fails */ ErrorCode InitLocalHotCache(); @@ -824,21 +824,6 @@ class Client { */ void UnregisterLocalHotCacheMemory(); - /** - * @brief Read MC_STORE_LOCAL_HOT_CACHE_SIZE from environment variable - * @return Cache size in bytes, or 0 if not set or invalid - */ - size_t GetLocalHotCacheSizeFromEnv(); - - /** - * @brief Read MC_STORE_LOCAL_HOT_BLOCK_SIZE from environment variable - * @param default_value Default block size to use if env var is not set or - * invalid - * @return Parsed block size from environment, or default_value if not - * set/invalid - */ - size_t GetLocalHotBlockSizeFromEnv(size_t default_value); - /** * @brief Redirect replica descriptor to local hot cache if cache hit * @param key Object key diff --git a/mooncake-store/include/local_hot_cache.h b/mooncake-store/include/local_hot_cache.h index 3fa827f331..4ecd6bd953 100644 --- a/mooncake-store/include/local_hot_cache.h +++ b/mooncake-store/include/local_hot_cache.h @@ -19,6 +19,15 @@ namespace mooncake { +struct LocalHotCacheConfig { + size_t total_size_bytes = 0; + size_t block_size_bytes = 16 * 1024 * 1024; + bool use_shm = false; + uint8_t admission_threshold = 2; + + static LocalHotCacheConfig FromEnvironment(); +}; + /** * @brief Token captured at async hot cache fill submission time. * Invalidated when RemoveHotKey, BumpKeyGeneration, or Clear bumps diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 4eb689df05..3add24c8eb 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -92,6 +92,7 @@ set(MOONCAKE_STORE_CLIENT_SOURCES file_storage.cpp config/file_storage_config.cpp config/client_auto_port_config.cpp + config/local_hot_cache_config.cpp device/accelerator_device.cpp device/accelerator_registry.cpp device/runtime_accelerator.cpp diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 77291f2913..4b076a4e12 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -4975,60 +4975,6 @@ tl::expected Client::GetPreferredReplica( return replica_list[0]; } -size_t Client::GetLocalHotCacheSizeFromEnv() { - if (const char* ev_size = std::getenv("MC_STORE_LOCAL_HOT_CACHE_SIZE")) { - std::string ev_size_str(ev_size); - std::string error_msg = "Invalid MC_STORE_LOCAL_HOT_CACHE_SIZE='" + - ev_size_str + "', disable local hot cache"; - // Check for negative values - if (!ev_size_str.empty() && ev_size_str[0] == '-') { - LOG(WARNING) << error_msg; - return 0; - } - try { - unsigned long long v = std::stoull(ev_size_str, nullptr, 10); - if (v > 0) { - return static_cast(v); - } else { - LOG(WARNING) << error_msg; - return 0; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - return 0; - } - } - return 0; -} - -size_t Client::GetLocalHotBlockSizeFromEnv(size_t default_value) { - if (const char* ev_block_size = - std::getenv("MC_STORE_LOCAL_HOT_BLOCK_SIZE")) { - std::string ev_block_size_str(ev_block_size); - std::string error_msg = "Invalid MC_STORE_LOCAL_HOT_BLOCK_SIZE='" + - ev_block_size_str + - "', using default block size"; - // Check for negative values - if (!ev_block_size_str.empty() && ev_block_size_str[0] == '-') { - LOG(WARNING) << error_msg; - return default_value; - } - try { - unsigned long long v = std::stoull(ev_block_size_str, nullptr, 10); - if (v > 0) { - return static_cast(v); - } else { - LOG(WARNING) << error_msg; - return default_value; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - return default_value; - } - } - return default_value; -} - ErrorCode Client::InitLocalHotCache() { hot_cache_handler_.reset(); UnregisterLocalHotCacheMemory(); @@ -5039,38 +4985,22 @@ ErrorCode Client::InitLocalHotCache() { return ErrorCode::OK; } - // Defaults: hot cache is disabled unless MC_STORE_LOCAL_HOT_CACHE_SIZE is - // set to a positive value; when enabled, default block size is 16MB and - // thread_num is 2. - size_t block_size = 16 * 1024 * 1024; // 16MB default block size - size_t thread_num = 2; - - // Read MC_STORE_LOCAL_HOT_CACHE_SIZE from environment - size_t total_cache = GetLocalHotCacheSizeFromEnv(); - if (total_cache == 0) { - // Environment variable not set or invalid, disable cache + const auto config = LocalHotCacheConfig::FromEnvironment(); + if (config.total_size_bytes == 0) { return ErrorCode::OK; } - // Read MC_STORE_LOCAL_HOT_BLOCK_SIZE from environment - block_size = GetLocalHotBlockSizeFromEnv(block_size); - - // MC_STORE_LOCAL_HOT_CACHE_USE_SHM: "1" enables memfd-backed shm (default - // off). When enabled, hot cache is shareable with dummy clients via IPC. - bool use_shm = false; - if (const char* ev = std::getenv("MC_STORE_LOCAL_HOT_CACHE_USE_SHM")) { - use_shm = (std::string(ev) == "1"); - } + size_t thread_num = 2; // Enable hot cache { - hot_cache_ = - std::make_shared(total_cache, block_size, use_shm); + hot_cache_ = std::make_shared( + config.total_size_bytes, config.block_size_bytes, config.use_shm); // Check if cache initialization was successful if (hot_cache_->GetCacheSize() == 0) { LOG(ERROR) << "Local hot cache creation failed: no blocks allocated. " - << "total_cache=" << total_cache; + << "total_cache=" << config.total_size_bytes; hot_cache_.reset(); hot_cache_handler_.reset(); admission_sketch_.reset(); @@ -5094,35 +5024,17 @@ ErrorCode Client::InitLocalHotCache() { } hot_cache_memory_registered_ = true; - LOG(INFO) << "Local hot cache enabled with cache size=" << total_cache - << ", block size=" << block_size + LOG(INFO) << "Local hot cache enabled with cache size=" + << config.total_size_bytes + << ", block size=" << config.block_size_bytes << ", block amount=" << hot_cache_->GetCacheSize() - << ", shm=" << (use_shm ? "on" : "off") + << ", shm=" << (config.use_shm ? "on" : "off") << ", transfer engine registered=on"; // Create async handler with 2 worker threads hot_cache_handler_ = std::make_unique(hot_cache_, thread_num); admission_sketch_ = std::make_unique(); - - // MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD: minimum CMS count before a - // key is admitted to hot cache (default 2). - if (const char* ev = - std::getenv("MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD")) { - std::string ev_str(ev); - std::string error_msg = - "Invalid MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD='" + ev_str + - "', using default"; - try { - unsigned long long v = std::stoull(ev_str, nullptr, 10); - if (v > 0 && v <= 255) { - admission_threshold_ = static_cast(v); - } else { - LOG(WARNING) << error_msg; - } - } catch (const std::exception&) { - LOG(WARNING) << error_msg; - } - } + admission_threshold_ = config.admission_threshold; } return ErrorCode::OK; } diff --git a/mooncake-store/src/config/local_hot_cache_config.cpp b/mooncake-store/src/config/local_hot_cache_config.cpp new file mode 100644 index 0000000000..ab77f00332 --- /dev/null +++ b/mooncake-store/src/config/local_hot_cache_config.cpp @@ -0,0 +1,92 @@ +#include "local_hot_cache.h" + +#include + +#include +#include +#include + +#include "environ.h" +#include "environment_variables.h" + +namespace mooncake { +namespace { + +size_t ParseLegacyPositiveSizeOr(const std::optional& raw_value, + size_t default_value, + const char* variable_name, + const char* fallback_message) { + if (!raw_value.has_value()) { + return default_value; + } + const std::string error_message = "Invalid " + std::string(variable_name) + + "='" + *raw_value + "'" + + fallback_message; + if (!raw_value->empty() && raw_value->front() == '-') { + LOG(WARNING) << error_message; + return default_value; + } + + try { + // Keep the legacy numeric-prefix behavior of std::stoull. + const unsigned long long value = std::stoull(*raw_value, nullptr, 10); + if (value > 0) { + return static_cast(value); + } + } catch (const std::exception&) { + } + + LOG(WARNING) << error_message; + return default_value; +} + +} // namespace + +LocalHotCacheConfig LocalHotCacheConfig::FromEnvironment() { + LocalHotCacheConfig config; + using Variables = LocalHotCacheEnvironmentVariables; + + const auto total_size = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_CACHE_SIZE); + config.total_size_bytes = ParseLegacyPositiveSizeOr( + total_size, 0, Variables::MC_STORE_LOCAL_HOT_CACHE_SIZE.name, + ", disable local hot cache"); + if (config.total_size_bytes == 0) { + return config; + } + + const auto block_size = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_BLOCK_SIZE); + config.block_size_bytes = + ParseLegacyPositiveSizeOr(block_size, config.block_size_bytes, + Variables::MC_STORE_LOCAL_HOT_BLOCK_SIZE.name, + ", using default block size"); + + config.use_shm = Environ::Read(Variables::MC_STORE_LOCAL_HOT_CACHE_USE_SHM) + .value_or(std::string{}) == "1"; + + const auto admission_threshold = + Environ::Read(Variables::MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD); + if (admission_threshold.has_value()) { + const std::string error_message = + "Invalid " + + std::string( + Variables::MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD.name) + + "='" + *admission_threshold + "', using default"; + try { + const unsigned long long value = + std::stoull(*admission_threshold, nullptr, 10); + if (value > 0 && value <= 255) { + config.admission_threshold = static_cast(value); + } else { + LOG(WARNING) << error_message; + } + } catch (const std::exception&) { + LOG(WARNING) << error_message; + } + } + + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index ed3c035f97..24c516096b 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -145,6 +145,7 @@ add_store_test(offset_allocator_test offset_allocator_test.cpp) add_store_test(utils_test utils_test.cpp) add_store_test(client_buffer_test client_buffer_test.cpp) add_store_test(client_auto_port_config_test client_auto_port_config_test.cpp) +add_store_test(local_hot_cache_config_test local_hot_cache_config_test.cpp) add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp) add_store_test(client_tcp_local_memcpy_test client_tcp_local_memcpy_test.cpp) add_store_test(pybind_client_test pybind_client_test.cpp) diff --git a/mooncake-store/tests/local_hot_cache_config_test.cpp b/mooncake-store/tests/local_hot_cache_config_test.cpp new file mode 100644 index 0000000000..354dfd4de0 --- /dev/null +++ b/mooncake-store/tests/local_hot_cache_config_test.cpp @@ -0,0 +1,150 @@ +#include +#include + +#include +#include +#include + +#include "local_hot_cache.h" + +namespace mooncake { +namespace { + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + if (const char* value = std::getenv(name)) { + original_ = value; + } + unsetenv(name); + } + + ~ScopedEnvVar() { + if (original_.has_value()) { + setenv(name_.c_str(), original_->c_str(), 1); + } else { + unsetenv(name_.c_str()); + } + } + + void Set(const char* value) { setenv(name_.c_str(), value, 1); } + + private: + std::string name_; + std::optional original_; +}; + +struct LocalHotCacheEnvironment { + ScopedEnvVar total_size{"MC_STORE_LOCAL_HOT_CACHE_SIZE"}; + ScopedEnvVar block_size{"MC_STORE_LOCAL_HOT_BLOCK_SIZE"}; + ScopedEnvVar use_shm{"MC_STORE_LOCAL_HOT_CACHE_USE_SHM"}; + ScopedEnvVar admission_threshold{"MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD"}; +}; + +class LocalHotCacheConfigTest : public ::testing::Test { + protected: + LocalHotCacheEnvironment env; + + void SetUp() override { + google::InitGoogleLogging("LocalHotCacheConfigTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } +}; + +TEST_F(LocalHotCacheConfigTest, UsesExistingDefaultsWhenEnvironmentIsUnset) { + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 0); + EXPECT_EQ(config.block_size_bytes, 16 * 1024 * 1024); + EXPECT_FALSE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 2); +} + +TEST_F(LocalHotCacheConfigTest, ReadsValidValues) { + env.total_size.Set("33554432"); + env.block_size.Set("4194304"); + env.use_shm.Set("1"); + env.admission_threshold.Set("5"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 32 * 1024 * 1024); + EXPECT_EQ(config.block_size_bytes, 4 * 1024 * 1024); + EXPECT_TRUE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 5); +} + +TEST_F(LocalHotCacheConfigTest, DisabledCacheDoesNotReadDependentSettings) { + env.total_size.Set("0"); + env.block_size.Set("4194304"); + env.use_shm.Set("1"); + env.admission_threshold.Set("5"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 0); + EXPECT_EQ(config.block_size_bytes, 16 * 1024 * 1024); + EXPECT_FALSE(config.use_shm); + EXPECT_EQ(config.admission_threshold, 2); +} + +TEST_F(LocalHotCacheConfigTest, InvalidCacheSizesDisableCache) { + for (const char* value : + {"", "0", "-1", "invalid", "18446744073709551616"}) { + env.total_size.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().total_size_bytes, 0) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, InvalidBlockSizesUseDefault) { + env.total_size.Set("33554432"); + + for (const char* value : + {"", "0", "-1", "invalid", "18446744073709551616"}) { + env.block_size.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().block_size_bytes, + 16 * 1024 * 1024) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, InvalidAdmissionThresholdsUseDefault) { + env.total_size.Set("33554432"); + + for (const char* value : + {"", "0", "-1", "256", "invalid", "18446744073709551616"}) { + env.admission_threshold.Set(value); + EXPECT_EQ(LocalHotCacheConfig::FromEnvironment().admission_threshold, 2) + << value; + } +} + +TEST_F(LocalHotCacheConfigTest, PreservesLegacyNumericPrefixParsing) { + env.total_size.Set("33554432suffix"); + env.block_size.Set("4194304suffix"); + env.admission_threshold.Set("5suffix"); + + const auto config = LocalHotCacheConfig::FromEnvironment(); + + EXPECT_EQ(config.total_size_bytes, 32 * 1024 * 1024); + EXPECT_EQ(config.block_size_bytes, 4 * 1024 * 1024); + EXPECT_EQ(config.admission_threshold, 5); +} + +TEST_F(LocalHotCacheConfigTest, SharedMemoryRequiresExactOne) { + env.total_size.Set("33554432"); + + for (const char* value : {"", "0", "true", "01", " 1"}) { + env.use_shm.Set(value); + EXPECT_FALSE(LocalHotCacheConfig::FromEnvironment().use_shm) << value; + } + + env.use_shm.Set("1"); + EXPECT_TRUE(LocalHotCacheConfig::FromEnvironment().use_shm); +} + +} // namespace +} // namespace mooncake