[Store] Structure RpcTimeoutConfig environment settings

mooncake-common/include/environment_variables.h

@@ -100,6 +100,13 @@ struct ReplicaSelectionEnvironmentVariables {
     MC_DEFINE_ENV_VAR(std::string, MC_STORE_REPLICA_SCORING);
 };
 
+struct RpcTimeoutEnvironmentVariables {
+    // Preserve atoll parsing: empty/nonnumeric values become zero, and numeric
+    // prefixes are accepted, unlike the shared typed integer parser.
+    MC_DEFINE_ENV_VAR(std::string, MC_RPC_TIMEOUT_MS);
+    MC_DEFINE_ENV_VAR(std::string, MC_RPC_CONNECT_TIMEOUT_MS);
+};
+
 #undef MC_DEFINE_ENV_VAR
 
 }  // namespace mooncake

mooncake-store/include/config/rpc_timeout_config.h

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <chrono>
+#include <optional>
+
+namespace mooncake {
+
+struct RpcTimeoutConfig {
+    // An absent override leaves the caller's connection policy unchanged.
+    std::optional<std::chrono::milliseconds> request_timeout;
+    std::optional<std::chrono::milliseconds> connect_timeout;
+
+    static RpcTimeoutConfig FromEnvironment();
+};
+
+}  // namespace mooncake

mooncake-store/include/master_client.h

@@ -13,6 +13,7 @@
 #include <ylt/coro_io/ibverbs/ib_socket.hpp>
 
 #include "client_metric.h"
+#include "config/rpc_timeout_config.h"
 #include "replica.h"
 #include "segment.h"
 #include "types.h"
@@ -48,17 +49,16 @@ inline void MaybeEnableRdmaSocketConfig(SocketConfigVariant& socket_config) {
 
 // Applies the RPC timeout overrides to any mooncake-store RPC client config.
 // Shared by the client->master pool and the store->store offload pool so the
-// two cannot drift. Unset variables leave coro_rpc's built-in 30s defaults in
-// place. A negative request timeout disables the per-request timer.
+// two cannot drift. Absent overrides leave the caller's defaults in place.
+// A negative request timeout disables the per-request timer.
 template <typename ClientConfig>
-inline void ApplyRpcTimeoutEnvOverrides(ClientConfig& client_config) {
-    if (const char* timeout_ms = std::getenv("MC_RPC_TIMEOUT_MS")) {
-        client_config.request_timeout_duration =
-            std::chrono::milliseconds(std::atoll(timeout_ms));
+inline void ApplyRpcTimeoutOverrides(ClientConfig& client_config,
+                                     const RpcTimeoutConfig& config) {
+    if (config.request_timeout.has_value()) {
+        client_config.request_timeout_duration = *config.request_timeout;
     }
-    if (const char* connect_ms = std::getenv("MC_RPC_CONNECT_TIMEOUT_MS")) {
-        client_config.connect_timeout_duration =
-            std::chrono::milliseconds(std::atoll(connect_ms));
+    if (config.connect_timeout.has_value()) {
+        client_config.connect_timeout_duration = *config.connect_timeout;
     }
 }
 
@@ -69,7 +69,8 @@ inline RpcClientPool::PoolConfig MakeMasterRpcClientPoolConfig() {
         MaybeEnableRdmaSocketConfig(config.client_config.socket_config);
     }
 
-    ApplyRpcTimeoutEnvOverrides(config.client_config);
+    ApplyRpcTimeoutOverrides(config.client_config,
+                             RpcTimeoutConfig::FromEnvironment());
     return config;
 }
 

mooncake-store/src/CMakeLists.txt

@@ -100,6 +100,7 @@ set(MOONCAKE_STORE_CLIENT_SOURCES
     config/local_hot_cache_config.cpp
     config/offset_allocator_backend_config.cpp
     config/replica_selection_config.cpp
+    config/rpc_timeout_config.cpp
     device/accelerator_device.cpp
     device/accelerator_registry.cpp
     device/runtime_accelerator.cpp

mooncake-store/src/config/rpc_timeout_config.cpp

@@ -0,0 +1,25 @@
+#include "config/rpc_timeout_config.h"
+
+#include <cstdlib>
+
+#include "environ.h"
+#include "environment_variables.h"
+
+namespace mooncake {
+
+RpcTimeoutConfig RpcTimeoutConfig::FromEnvironment() {
+    RpcTimeoutConfig config;
+    using Variables = RpcTimeoutEnvironmentVariables;
+    if (const auto value = Environ::Read(Variables::MC_RPC_TIMEOUT_MS)) {
+        config.request_timeout =
+            std::chrono::milliseconds{std::atoll(value->c_str())};
+    }
+    if (const auto value =
+            Environ::Read(Variables::MC_RPC_CONNECT_TIMEOUT_MS)) {
+        config.connect_timeout =
+            std::chrono::milliseconds{std::atoll(value->c_str())};
+    }
+    return config;
+}
+
+}  // namespace mooncake

mooncake-store/src/real_client.cpp

@@ -6991,7 +6991,8 @@ ClientRequester::ClientRequester() {
     // peer which is gone blocks for connect_retry_count * 30s plus the waits
     // between retries -- 91s with the defaults above -- and no configuration
     // can shorten it. Defaults are unchanged when the variables are unset.
-    detail::ApplyRpcTimeoutEnvOverrides(pool_conf.client_config);
+    detail::ApplyRpcTimeoutOverrides(pool_conf.client_config,
+                                     RpcTimeoutConfig::FromEnvironment());
 
     client_pools_ =
         std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(

mooncake-store/tests/CMakeLists.txt

@@ -126,6 +126,7 @@ add_test(
 set_tests_properties(object_checksum_client_test
                      PROPERTIES ENVIRONMENT "MOONCAKE_STORE_CHECKSUM=1")
 add_store_test(rpc_timeout_test rpc_timeout_test.cpp)
+add_store_test(rpc_timeout_config_test rpc_timeout_config_test.cpp)
 if(USE_CXL)
   add_store_test(cxl_client_integration_test cxl_client_integration_test.cpp)
 endif()

mooncake-store/tests/rpc_timeout_config_test.cpp

@@ -0,0 +1,106 @@
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+#include <optional>
+#include <string>
+
+#include "config/rpc_timeout_config.h"
+
+namespace mooncake {
+namespace {
+
+class RpcTimeoutConfigTest : public ::testing::Test {
+   protected:
+    void SetUp() override {
+        for (int i = 0; i < 2; ++i) {
+            if (const char* value = std::getenv(names_[i])) {
+                original_[i] = value;
+            }
+        }
+        for (const char* name : names_) {
+            ASSERT_EQ(unsetenv(name), 0);
+        }
+    }
+
+    void TearDown() override {
+        for (int i = 0; i < 2; ++i) {
+            if (original_[i].has_value()) {
+                EXPECT_EQ(setenv(names_[i], original_[i]->c_str(), 1), 0);
+            } else {
+                EXPECT_EQ(unsetenv(names_[i]), 0);
+            }
+        }
+    }
+
+    static constexpr const char* names_[] = {"MC_RPC_TIMEOUT_MS",
+                                             "MC_RPC_CONNECT_TIMEOUT_MS"};
+    std::optional<std::string> original_[2];
+};
+
+TEST_F(RpcTimeoutConfigTest, UnsetVariablesLeaveOverridesAbsent) {
+    const auto config = RpcTimeoutConfig::FromEnvironment();
+    EXPECT_FALSE(config.request_timeout.has_value());
+    EXPECT_FALSE(config.connect_timeout.has_value());
+}
+
+TEST_F(RpcTimeoutConfigTest, PreservesLegacyConversionForEachVariable) {
+    struct Case {
+        const char* value;
+        int64_t expected;
+    };
+    const Case cases[] = {
+        {"", 0},
+        {"garbage", 0},
+        {"0", 0},
+        {"1500", 1500},
+        {"-1", -1},
+        {" +42 ", 42},
+        {" \t", 0},
+        {"1500ms", 1500},
+        {"0x10", 0},
+        {"9223372036854775807", std::numeric_limits<int64_t>::max()},
+        {"-9223372036854775808", std::numeric_limits<int64_t>::min()},
+    };
+    for (int i = 0; i < 2; ++i) {
+        SCOPED_TRACE(names_[i]);
+        for (const auto& entry : cases) {
+            SCOPED_TRACE(entry.value);
+            ASSERT_EQ(setenv(names_[i], entry.value, 1), 0);
+            const auto config = RpcTimeoutConfig::FromEnvironment();
+            const auto& selected =
+                i == 0 ? config.request_timeout : config.connect_timeout;
+            const auto& other =
+                i == 0 ? config.connect_timeout : config.request_timeout;
+            ASSERT_TRUE(selected.has_value());
+            EXPECT_EQ(*selected, std::chrono::milliseconds(entry.expected));
+            EXPECT_FALSE(other.has_value());
+        }
+        ASSERT_EQ(unsetenv(names_[i]), 0);
+    }
+}
+
+TEST_F(RpcTimeoutConfigTest, EachConstructionReadsCurrentEnvironment) {
+    ASSERT_EQ(setenv("MC_RPC_TIMEOUT_MS", "1500", 1), 0);
+    ASSERT_EQ(setenv("MC_RPC_CONNECT_TIMEOUT_MS", "1000", 1), 0);
+    const auto original = RpcTimeoutConfig::FromEnvironment();
+
+    ASSERT_EQ(setenv("MC_RPC_TIMEOUT_MS", "0", 1), 0);
+    ASSERT_EQ(setenv("MC_RPC_CONNECT_TIMEOUT_MS", "-1", 1), 0);
+    const auto changed = RpcTimeoutConfig::FromEnvironment();
+    EXPECT_EQ(changed.request_timeout, std::chrono::milliseconds(0));
+    EXPECT_EQ(changed.connect_timeout, std::chrono::milliseconds(-1));
+
+    ASSERT_EQ(unsetenv("MC_RPC_TIMEOUT_MS"), 0);
+    ASSERT_EQ(unsetenv("MC_RPC_CONNECT_TIMEOUT_MS"), 0);
+    const auto cleared = RpcTimeoutConfig::FromEnvironment();
+    EXPECT_FALSE(cleared.request_timeout.has_value());
+    EXPECT_FALSE(cleared.connect_timeout.has_value());
+    EXPECT_EQ(original.request_timeout, std::chrono::milliseconds(1500));
+    EXPECT_EQ(original.connect_timeout, std::chrono::milliseconds(1000));
+}
+
+}  // namespace
+}  // namespace mooncake