diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 2eb5352f0e..8d2d337ef7 100644 --- a/mooncake-common/include/environment_variables.h +++ b/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); +}; + struct LocalFileSnapshotEnvironmentVariables { MC_DEFINE_ENV_VAR(std::string, MOONCAKE_SNAPSHOT_LOCAL_PATH); }; diff --git a/mooncake-store/include/config/rpc_timeout_config.h b/mooncake-store/include/config/rpc_timeout_config.h new file mode 100644 index 0000000000..3b20496cce --- /dev/null +++ b/mooncake-store/include/config/rpc_timeout_config.h @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace mooncake { + +struct RpcTimeoutConfig { + // An absent override leaves the caller's connection policy unchanged. + std::optional request_timeout; + std::optional connect_timeout; + + static RpcTimeoutConfig FromEnvironment(); +}; + +} // namespace mooncake diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 80c9bac4b2..33b3e87058 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -13,6 +13,7 @@ #include #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 -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; } diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 3b2e06bcf8..daf01cb620 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -101,6 +101,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 config/transfer_submitter_config.cpp device/accelerator_device.cpp device/accelerator_registry.cpp diff --git a/mooncake-store/src/config/rpc_timeout_config.cpp b/mooncake-store/src/config/rpc_timeout_config.cpp new file mode 100644 index 0000000000..65777301be --- /dev/null +++ b/mooncake-store/src/config/rpc_timeout_config.cpp @@ -0,0 +1,25 @@ +#include "config/rpc_timeout_config.h" + +#include + +#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 diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 229dfd45f0..d82ac4b09a 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -6989,7 +6989,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>( diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index fae1180849..8fb8dc1db4 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -138,6 +138,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() diff --git a/mooncake-store/tests/rpc_timeout_config_test.cpp b/mooncake-store/tests/rpc_timeout_config_test.cpp new file mode 100644 index 0000000000..efb5cf5591 --- /dev/null +++ b/mooncake-store/tests/rpc_timeout_config_test.cpp @@ -0,0 +1,106 @@ +#include + +#include +#include +#include +#include +#include +#include + +#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 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::max()}, + {"-9223372036854775808", std::numeric_limits::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 diff --git a/mooncake-store/tests/rpc_timeout_test.cpp b/mooncake-store/tests/rpc_timeout_test.cpp index b2e31f46de..5ab4adc1cb 100644 --- a/mooncake-store/tests/rpc_timeout_test.cpp +++ b/mooncake-store/tests/rpc_timeout_test.cpp @@ -1,10 +1,11 @@ -// Unit tests for the client->master RPC timeout feature. +// Tests for the client->master and store->store RPC timeout overrides. // // Covers: // 1. The ErrorCode::RPC_TIMEOUT enum value and its toString() mapping. // 2. End-to-end: MC_RPC_TIMEOUT_MS shortens the per-request deadline so that // an unresponsive master surfaces ErrorCode::RPC_TIMEOUT (not RPC_FAIL), // and it does so within the configured budget rather than the 30s default. +// 3. The offload requester honors the same timeout and error mapping. // // The end-to-end test points a MasterClient at a "black hole" TCP listener: a // socket that accepts the connection (so connect() succeeds) but never sends a @@ -21,14 +22,45 @@ #include #include +#include #include #include "master_client.h" +#include "pyclient.h" #include "types.h" namespace mooncake { namespace { +class RpcTimeoutEnvTest : public ::testing::Test { + protected: + void SetUp() override { + for (int i = 0; i < 3; ++i) { + if (const char* value = std::getenv(names_[i])) { + original_[i] = value; + } + } + // These integration tests use TCP regardless of the ambient protocol. + for (const char* name : names_) { + ASSERT_EQ(unsetenv(name), 0); + } + } + + void TearDown() override { + for (int i = 0; i < 3; ++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", "MC_RPC_PROTOCOL"}; + std::optional original_[3]; +}; + // A TCP listener on 127.0.0.1 that accepts connections but never replies. // The kernel completes the TCP handshake for backlogged connections, so a // client's connect() succeeds; any subsequent request simply never gets a @@ -84,7 +116,7 @@ TEST(RpcTimeoutTest, ErrorCodeValueAndString) { // End-to-end: a small MC_RPC_TIMEOUT_MS must be honored by every RPC and turn // an unanswered call into ErrorCode::RPC_TIMEOUT well before the 30s default. -TEST(RpcTimeoutTest, RpcTimesOutAgainstUnresponsiveMaster) { +TEST_F(RpcTimeoutEnvTest, RpcTimesOutAgainstUnresponsiveMaster) { constexpr int kTimeoutMs = 500; // The constructor reads MC_RPC_TIMEOUT_MS, so it must be set beforehand. @@ -102,8 +134,6 @@ TEST(RpcTimeoutTest, RpcTimesOutAgainstUnresponsiveMaster) { std::chrono::steady_clock::now() - start) .count(); - ::unsetenv("MC_RPC_TIMEOUT_MS"); - // The unanswered ServiceReady RPC must be reported as a timeout, not as a // generic RPC failure. EXPECT_EQ(rc, ErrorCode::RPC_TIMEOUT) @@ -124,7 +154,7 @@ TEST(RpcTimeoutTest, RpcTimesOutAgainstUnresponsiveMaster) { // between retries with no way to configure it down. Both pools now go through // one helper, so cover the helper itself: an unset variable must leave the // built-in default alone, and a set one must be applied. -TEST(RpcTimeoutTest, TimeoutEnvOverridesAreOptIn) { +TEST_F(RpcTimeoutEnvTest, TimeoutEnvOverridesAreOptIn) { struct StubClientConfig { std::chrono::milliseconds request_timeout_duration{ std::chrono::seconds(30)}; @@ -132,20 +162,20 @@ TEST(RpcTimeoutTest, TimeoutEnvOverridesAreOptIn) { std::chrono::seconds(30)}; }; - ::unsetenv("MC_RPC_TIMEOUT_MS"); - ::unsetenv("MC_RPC_CONNECT_TIMEOUT_MS"); - StubClientConfig defaults; - detail::ApplyRpcTimeoutEnvOverrides(defaults); + detail::ApplyRpcTimeoutOverrides(defaults, + RpcTimeoutConfig::FromEnvironment()); EXPECT_EQ(defaults.request_timeout_duration, std::chrono::seconds(30)); EXPECT_EQ(defaults.connect_timeout_duration, std::chrono::seconds(30)); + const auto original_master_config = detail::MakeMasterRpcClientPoolConfig(); ASSERT_EQ(::setenv("MC_RPC_TIMEOUT_MS", "1500", /*overwrite=*/1), 0); ASSERT_EQ(::setenv("MC_RPC_CONNECT_TIMEOUT_MS", "1000", /*overwrite=*/1), 0); StubClientConfig overridden; - detail::ApplyRpcTimeoutEnvOverrides(overridden); + detail::ApplyRpcTimeoutOverrides(overridden, + RpcTimeoutConfig::FromEnvironment()); EXPECT_EQ(overridden.request_timeout_duration, std::chrono::milliseconds(1500)); EXPECT_EQ(overridden.connect_timeout_duration, @@ -153,11 +183,71 @@ TEST(RpcTimeoutTest, TimeoutEnvOverridesAreOptIn) { // The master pool is built from the same helper, so it sees them too. auto master_config = detail::MakeMasterRpcClientPoolConfig(); + EXPECT_EQ(master_config.client_config.request_timeout_duration, + std::chrono::milliseconds(1500)); EXPECT_EQ(master_config.client_config.connect_timeout_duration, std::chrono::milliseconds(1000)); + EXPECT_EQ(original_master_config.client_config.request_timeout_duration, + std::chrono::seconds(30)); + EXPECT_EQ(original_master_config.client_config.connect_timeout_duration, + std::chrono::seconds(30)); +} + +TEST_F(RpcTimeoutEnvTest, OverridesPreserveCallerPolicyAndUseResolvedValues) { + struct StubClientConfig { + std::chrono::milliseconds request_timeout_duration{7000}; + std::chrono::milliseconds connect_timeout_duration{1000}; + }; + StubClientConfig defaults; + detail::ApplyRpcTimeoutOverrides(defaults, + RpcTimeoutConfig::FromEnvironment()); + EXPECT_EQ(defaults.request_timeout_duration, + std::chrono::milliseconds(7000)); + EXPECT_EQ(defaults.connect_timeout_duration, + std::chrono::milliseconds(1000)); + + // The helper must apply this snapshot, not re-read the environment. + ASSERT_EQ(setenv("MC_RPC_TIMEOUT_MS", "9999", 1), 0); + ASSERT_EQ(setenv("MC_RPC_CONNECT_TIMEOUT_MS", "9999", 1), 0); + for (const int timeout_ms : {0, -1, 1500}) { + SCOPED_TRACE(timeout_ms); + RpcTimeoutConfig config; + config.request_timeout = std::chrono::milliseconds(timeout_ms); + StubClientConfig request_only; + detail::ApplyRpcTimeoutOverrides(request_only, config); + EXPECT_EQ(request_only.request_timeout_duration, + std::chrono::milliseconds(timeout_ms)); + EXPECT_EQ(request_only.connect_timeout_duration, + std::chrono::milliseconds(1000)); + + config.request_timeout.reset(); + config.connect_timeout = std::chrono::milliseconds(timeout_ms); + StubClientConfig connect_only; + detail::ApplyRpcTimeoutOverrides(connect_only, config); + EXPECT_EQ(connect_only.request_timeout_duration, + std::chrono::milliseconds(7000)); + EXPECT_EQ(connect_only.connect_timeout_duration, + std::chrono::milliseconds(timeout_ms)); + } +} + +TEST_F(RpcTimeoutEnvTest, RpcTimesOutAgainstUnresponsiveOffloadPeer) { + constexpr int kTimeoutMs = 500; + ASSERT_EQ(setenv("MC_RPC_TIMEOUT_MS", "500", 1), 0); + BlackHoleServer server; + ClientRequester requester; + + const auto start = std::chrono::steady_clock::now(); + const auto result = + requester.batch_get_offload_object(server.address(), {"key"}, {1}); + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); - ::unsetenv("MC_RPC_TIMEOUT_MS"); - ::unsetenv("MC_RPC_CONNECT_TIMEOUT_MS"); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::RPC_TIMEOUT); + EXPECT_GE(elapsed, kTimeoutMs - 100); + EXPECT_LT(elapsed, 5000); } } // namespace mooncake