From 4a04c2dfb132fd82a004b77b7f42fae692e556d2 Mon Sep 17 00:00:00 2001 From: Schatten Date: Tue, 8 Sep 2026 18:22:06 +0800 Subject: [PATCH] [Store] Structure NoFRegisterConfig environment settings Signed-off-by: Schatten --- .../include/environment_variables.h | 5 + mooncake-store/src/CMakeLists.txt | 1 + .../src/config/nof_register_config.cpp | 30 +++++ .../src/config/nof_register_config.h | 13 ++ mooncake-store/src/ssd_register_client.cpp | 43 ++----- mooncake-store/tests/CMakeLists.txt | 1 + .../tests/nof_register_config_test.cpp | 112 ++++++++++++++++++ 7 files changed, 174 insertions(+), 31 deletions(-) create mode 100644 mooncake-store/src/config/nof_register_config.cpp create mode 100644 mooncake-store/src/config/nof_register_config.h create mode 100644 mooncake-store/tests/nof_register_config_test.cpp diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 50154c13c6..da4af34608 100644 --- a/mooncake-common/include/environment_variables.h +++ b/mooncake-common/include/environment_variables.h @@ -123,6 +123,11 @@ struct TransferSubmitterEnvironmentVariables { MC_DEFINE_ENV_VAR(std::string, MC_STORE_MEMCPY); }; +struct NoFRegisterEnvironmentVariables { + // Keep the raw string to preserve case normalization and warning behavior. + MC_DEFINE_ENV_VAR(std::string, MC_NOF_TRTYPE); +}; + #undef MC_DEFINE_ENV_VAR } // namespace mooncake diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 435206082d..cb769c4dd0 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -100,6 +100,7 @@ set(MOONCAKE_STORE_CLIENT_SOURCES config/file_per_key_config.cpp config/client_auto_port_config.cpp config/local_hot_cache_config.cpp + config/nof_register_config.cpp config/offset_allocator_backend_config.cpp config/replica_selection_config.cpp config/rpc_timeout_config.cpp diff --git a/mooncake-store/src/config/nof_register_config.cpp b/mooncake-store/src/config/nof_register_config.cpp new file mode 100644 index 0000000000..284298e6b5 --- /dev/null +++ b/mooncake-store/src/config/nof_register_config.cpp @@ -0,0 +1,30 @@ +#include "nof_register_config.h" + +#include + +#include +#include + +#include "environ.h" +#include "environment_variables.h" + +namespace mooncake { + +NoFRegisterConfig NoFRegisterConfig::FromEnvironment() { + NoFRegisterConfig config; + config.transport_type = + Environ::Read(NoFRegisterEnvironmentVariables::MC_NOF_TRTYPE) + .value_or("RDMA"); + std::transform(config.transport_type.begin(), config.transport_type.end(), + config.transport_type.begin(), [](unsigned char c) { + return static_cast(std::toupper(c)); + }); + if (config.transport_type != "RDMA" && config.transport_type != "TCP") { + LOG(WARNING) << "Invalid MC_NOF_TRTYPE=" << config.transport_type + << ", fallback to RDMA"; + config.transport_type = "RDMA"; + } + return config; +} + +} // namespace mooncake diff --git a/mooncake-store/src/config/nof_register_config.h b/mooncake-store/src/config/nof_register_config.h new file mode 100644 index 0000000000..93758ac0fb --- /dev/null +++ b/mooncake-store/src/config/nof_register_config.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace mooncake { + +struct NoFRegisterConfig { + std::string transport_type = "RDMA"; + + static NoFRegisterConfig FromEnvironment(); +}; + +} // namespace mooncake diff --git a/mooncake-store/src/ssd_register_client.cpp b/mooncake-store/src/ssd_register_client.cpp index 39cc9f393e..8bcd4bde4c 100644 --- a/mooncake-store/src/ssd_register_client.cpp +++ b/mooncake-store/src/ssd_register_client.cpp @@ -1,7 +1,6 @@ #include "ssd_register_client.h" -#include -#include -#include + +#include "config/nof_register_config.h" namespace mooncake { @@ -25,21 +24,12 @@ int NoFRegisterClient::set_register(const std::string &nqn, size_t nsid, return OPERATION_FAILED; } - const char *trtype_env = std::getenv("MC_NOF_TRTYPE"); - std::string trtype = trtype_env ? trtype_env : "RDMA"; - std::transform( - trtype.begin(), trtype.end(), trtype.begin(), - [](unsigned char c) { return static_cast(std::toupper(c)); }); - if (trtype != "RDMA" && trtype != "TCP") { - LOG(WARNING) << "Invalid MC_NOF_TRTYPE=" << trtype - << ", fallback to RDMA"; - trtype = "RDMA"; - } + const auto config = NoFRegisterConfig::FromEnvironment(); - std::string te_endpoint = "traddr:" + traddr + - " trsvcid:" + std::to_string(trsvcid) + - " subnqn:" + nqn + " trtype:" + trtype + - " adrfam:IPv4 ns:" + std::to_string(nsid); + std::string te_endpoint = + "traddr:" + traddr + " trsvcid:" + std::to_string(trsvcid) + + " subnqn:" + nqn + " trtype:" + config.transport_type + + " adrfam:IPv4 ns:" + std::to_string(nsid); NoFSegment segment; segment.base = base; @@ -70,22 +60,13 @@ int NoFRegisterClient::set_unregister_by_endpoint( return OPERATION_FAILED; } - const char *trtype_env = std::getenv("MC_NOF_TRTYPE"); - std::string trtype = trtype_env ? trtype_env : "RDMA"; - std::transform( - trtype.begin(), trtype.end(), trtype.begin(), - [](unsigned char c) { return static_cast(std::toupper(c)); }); - if (trtype != "RDMA" && trtype != "TCP") { - LOG(WARNING) << "Invalid MC_NOF_TRTYPE=" << trtype - << ", fallback to RDMA"; - trtype = "RDMA"; - } + const auto config = NoFRegisterConfig::FromEnvironment(); // Build the te_endpoint string to match registered segments - std::string te_endpoint = "traddr:" + traddr + - " trsvcid:" + std::to_string(trsvcid) + - " subnqn:" + nqn + " trtype:" + trtype + - " adrfam:IPv4 ns:" + std::to_string(nsid); + std::string te_endpoint = + "traddr:" + traddr + " trsvcid:" + std::to_string(trsvcid) + + " subnqn:" + nqn + " trtype:" + config.transport_type + + " adrfam:IPv4 ns:" + std::to_string(nsid); LOG(INFO) << "Built te_endpoint: " << te_endpoint; diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index b269084e79..6bf7961374 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -128,6 +128,7 @@ add_store_test(master_service_promotion_test_for_snapshot if(USE_NOF) add_store_test(nof_heartbeat_test nof_heartbeat_test.cpp) endif() +add_store_test(nof_register_config_test nof_register_config_test.cpp) add_store_test(client_integration_test client_integration_test.cpp) add_test( NAME object_checksum_client_test diff --git a/mooncake-store/tests/nof_register_config_test.cpp b/mooncake-store/tests/nof_register_config_test.cpp new file mode 100644 index 0000000000..1204ae4d4c --- /dev/null +++ b/mooncake-store/tests/nof_register_config_test.cpp @@ -0,0 +1,112 @@ +#include "../src/config/nof_register_config.h" + +#include +#include + +#include +#include +#include + +namespace mooncake::test { +namespace { + +class NoFRegisterConfigTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("NoFRegisterConfigTest"); + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + void SetUp() override { + original_logtostderr_ = FLAGS_logtostderr; + FLAGS_logtostderr = true; + if (const char* value = std::getenv("MC_NOF_TRTYPE")) { + original_ = value; + } + ASSERT_EQ(unsetenv("MC_NOF_TRTYPE"), 0); + } + + void TearDown() override { + if (original_.has_value()) { + EXPECT_EQ(setenv("MC_NOF_TRTYPE", original_->c_str(), 1), 0); + } else { + EXPECT_EQ(unsetenv("MC_NOF_TRTYPE"), 0); + } + FLAGS_logtostderr = original_logtostderr_; + } + + private: + std::optional original_; + bool original_logtostderr_ = false; +}; + +TEST_F(NoFRegisterConfigTest, UnsetTransportDefaultsToRdmaSilently) { + testing::internal::CaptureStderr(); + const auto config = NoFRegisterConfig::FromEnvironment(); + const auto diagnostics = testing::internal::GetCapturedStderr(); + + EXPECT_EQ(config.transport_type, "RDMA"); + EXPECT_TRUE(diagnostics.empty()) << diagnostics; +} + +TEST_F(NoFRegisterConfigTest, ValidTransportIsCaseInsensitive) { + struct Case { + const char* value; + const char* expected; + }; + const Case cases[] = {{"RDMA", "RDMA"}, {"rdma", "RDMA"}, {"RdMa", "RDMA"}, + {"TCP", "TCP"}, {"tcp", "TCP"}, {"TcP", "TCP"}}; + for (const auto& entry : cases) { + SCOPED_TRACE(entry.value); + ASSERT_EQ(setenv("MC_NOF_TRTYPE", entry.value, 1), 0); + + testing::internal::CaptureStderr(); + const auto config = NoFRegisterConfig::FromEnvironment(); + const auto diagnostics = testing::internal::GetCapturedStderr(); + + EXPECT_EQ(config.transport_type, entry.expected); + EXPECT_TRUE(diagnostics.empty()) << diagnostics; + } +} + +TEST_F(NoFRegisterConfigTest, EmptyAndInvalidValuesWarnAndFallbackToRdma) { + struct Case { + const char* value; + const char* normalized; + }; + const Case cases[] = {{"", ""}, + {"udp", "UDP"}, + {" tcp", " TCP"}, + {"tcp ", "TCP "}, + {"1", "1"}}; + for (const auto& entry : cases) { + SCOPED_TRACE(entry.value); + ASSERT_EQ(setenv("MC_NOF_TRTYPE", entry.value, 1), 0); + + testing::internal::CaptureStderr(); + const auto config = NoFRegisterConfig::FromEnvironment(); + const auto diagnostics = testing::internal::GetCapturedStderr(); + + EXPECT_EQ(config.transport_type, "RDMA"); + EXPECT_NE(diagnostics.find(std::string("Invalid MC_NOF_TRTYPE=") + + entry.normalized + ", fallback to RDMA"), + std::string::npos); + } +} + +TEST_F(NoFRegisterConfigTest, NewConfigsReadCurrentEnvironment) { + ASSERT_EQ(setenv("MC_NOF_TRTYPE", "tcp", 1), 0); + const auto first = NoFRegisterConfig::FromEnvironment(); + ASSERT_EQ(setenv("MC_NOF_TRTYPE", "rdma", 1), 0); + const auto second = NoFRegisterConfig::FromEnvironment(); + ASSERT_EQ(unsetenv("MC_NOF_TRTYPE"), 0); + const auto third = NoFRegisterConfig::FromEnvironment(); + + EXPECT_EQ(first.transport_type, "TCP"); + EXPECT_EQ(second.transport_type, "RDMA"); + EXPECT_EQ(third.transport_type, "RDMA"); +} + +} // namespace +} // namespace mooncake::test