From 5d77d3446b9fee89df1ea50468018b1315638a37 Mon Sep 17 00:00:00 2001 From: Xingyuan Wu Date: Fri, 16 Oct 2026 16:17:37 +0800 Subject: [PATCH 01/14] [Store] Add LastHitOnly lease selection to BatchProbeKey (#4314) * [Store] Add LastHitOnly lease selection to BatchProbeKey Select and lease only the last complete request-scoped candidate, returning a selected-only mask. Preserve lease-free default probes and existing existence APIs. Thread the policy through both RPC boundaries and Python, with lease, tenant, shared-group and concurrency coverage. Batch-probe RPC signatures change; participating components must upgrade together. Implements the Store portion of RFC #4255. * [Store] Use string_view for probe validation errors --- .../api-reference/python/mooncake-store.md | 69 ++++- mooncake-integration/store/store_py.cpp | 30 +- mooncake-store/include/client_service.h | 9 +- mooncake-store/include/dummy_client.h | 3 +- mooncake-store/include/master_client.h | 9 +- mooncake-store/include/master_service.h | 3 +- mooncake-store/include/pyclient.h | 3 +- mooncake-store/include/real_client.h | 13 +- mooncake-store/include/rpc_service.h | 2 +- mooncake-store/include/types.h | 9 + mooncake-store/src/client_service.cpp | 4 +- mooncake-store/src/dummy_client.cpp | 6 +- mooncake-store/src/master_client.cpp | 5 +- mooncake-store/src/master_service.cpp | 51 +++- mooncake-store/src/real_client.cpp | 10 +- mooncake-store/src/rpc_service.cpp | 6 +- mooncake-store/tests/CMakeLists.txt | 2 + .../tests/dummy_client_probe_key_test.cpp | 19 ++ .../master_service_probe_policy_test.cpp | 284 ++++++++++++++++++ mooncake-store/tests/master_service_test.cpp | 2 +- mooncake-store/tests/pybind_client_test.cpp | 84 ++++++ .../tests/test_batch_probe_policy.py | 85 ++++++ 22 files changed, 659 insertions(+), 49 deletions(-) create mode 100644 mooncake-store/tests/master_service_probe_policy_test.cpp create mode 100644 mooncake-wheel/tests/test_batch_probe_policy.py diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md index 42a1419937..b860821ca8 100644 --- a/docs/source/api-reference/python/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -1619,19 +1619,51 @@ else: --- #### batch_probe_key() -Point-in-time existence check for multiple objects in a single batch -operation, granting no read leases. +Probe multiple objects, optionally selecting and leasing the last complete +candidate. The default `"none"` policy remains a lease-free existence check. ```python -def batch_probe_key(self, keys: List[str]) -> List[int] +def batch_probe_key( + self, keys: List[str], policy: str = "none", candidate_size: int = 1 +) -> List[int] ``` **Parameters:** - `keys` (List[str]): List of object identifiers to check - -**Returns:** -- `List[int]`: List of existence results (1=existed at probe time, -0=not exists, -1=error) +- `policy` (str): `"none"` or `"LastHitOnly"`. +- `candidate_size` (int): Number of consecutive keys forming each candidate. + Ignored for `"none"`. For a nonempty `"LastHitOnly"` request, it must be + positive and divide the number of keys exactly. Empty input returns `[]`. + +**Returns:** +- Both policies return per-key existence results in the same length and order + as `keys`: 1=readable when checked, 0=missing or unreadable. Negative values + indicate errors, not cache misses. +- With `"LastHitOnly"`, the last group of `candidate_size` all-1 results is + the selected, leased candidate. Earlier complete groups and incomplete + groups retain their existence results. If no group is all-1, no complete + candidate was selected. + +Candidates are ordered by the caller. A candidate is complete when every key +has a readable replica. Earlier candidates may be missing or incomplete. +The caller must include all required components, +cache groups, and rank shards; Store does not infer prefix or model semantics. +Duplicate key positions are allowed and must not alter the grouping. + +The Master first probes without leases, then rechecks and leases complete +candidates in reverse order using the existing per-shard lookup. Recheck +results replace that candidate's initial results; a candidate that loses a +member cannot remain all-1. A failed candidate may retain partial leases +while selection falls back to an earlier candidate. This is not an atomic +snapshot of all keys, and unselected existence results are not lease guarantees. + +Selected objects receive the Master's default read-lease TTL, protecting them +from normal eviction while the leases remain valid; soft-pin deadlines are +unchanged. The policy does not protect against lease expiry, forced removal, +or replica loss. Existing shared group leases (or a physical key repeated +outside the selected candidate) may also protect unselected positions. No +persistent group registration is added. Invalid candidate sizes return +per-key parameter errors; unknown policy strings raise `ValueError` in Python. **Example:** ```python @@ -1641,6 +1673,29 @@ candidates = [key for key, exists in zip(keys, results) if exists == 1] print("Probed candidates (unprotected from eviction):", candidates) ``` +For checkpoint selection, first use `batch_is_exist()` to protect the required +KV data and determine its usable boundary. Pass only valid checkpoint +candidates within that boundary, in increasing resume-position order: + +```python +keys = ["checkpoint128.rank0", "checkpoint128.rank1", + "checkpoint256.rank0", "checkpoint256.rank1"] +results = store.batch_probe_key(keys, policy="LastHitOnly", candidate_size=2) +if any(result < 0 for result in results): + raise RuntimeError("Checkpoint lookup failed") +selected = next((i for i in reversed(range(len(keys) // 2)) + if all(result == 1 for result in results[i * 2:(i + 1) * 2])), None) +# Both complete: [1, 1, 1, 1] selects 1. +# Last incomplete: [1, 1, 1, 0] selects 0. +# Load only the selected checkpoint, with the KV prefix it needs. +``` + +**Upgrade note:** the policy changes the `BatchProbeKey` Master RPC and the +DummyClient-to-RealClient batch-probe RPC signatures. Upgrade the Master, +client daemon, and clients together; mixed-version batch-probe calls are not +supported. Default Python arguments preserve the old call syntax, not wire +compatibility. `is_exist()`, `batch_is_exist()`, and `probe_key()` are unchanged. + --- #### get_size() diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 2dff85bd55..9080a079f2 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -2625,15 +2625,31 @@ PYBIND11_MODULE(_store, m) { .def( "batch_probe_key", [](MooncakeStorePyWrapper &self, - const std::vector &keys) { + const std::vector &keys, const std::string &policy, + uint64_t candidate_size) { + GrantLeasePolicy lease_policy; + if (policy == "none") { + lease_policy.lease_mode = ProbeLeaseMode::None; + } else if (policy == "LastHitOnly") { + lease_policy.lease_mode = ProbeLeaseMode::LastHitOnly; + } else { + throw py::value_error("Unknown batch_probe_key policy: " + + policy); + } + lease_policy.candidate_size = candidate_size; + if (!self.is_client_initialized()) { + return std::vector( + keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); + } py::gil_scoped_release release; - return self.store_->batchProbeKey(keys); + return self.store_->batchProbeKey(keys, lease_policy); }, - py::arg("keys"), - "Point-in-time existence check for multiple objects that grants " - "no read leases. Returns list of results: 1 if existed at the " - "time of the call, 0 if not exists, -1 if error. Objects may " - "still be evicted before a subsequent get.") + py::arg("keys"), py::arg("policy") = "none", + py::arg("candidate_size") = 1, + "Per-key existence results: 1 if readable, 0 if missing, and " + "negative values on error. none grants no leases. LastHitOnly " + "leases the last all-1 group of candidate_size results; failed " + "candidates may retain partial leases after a concurrent eviction.") .def( "close", [](MooncakeStorePyWrapper &self) { diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index d072a36686..94e203b5b3 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -468,13 +468,14 @@ class Client { const std::vector& keys); /** - * @brief Point-in-time existence check for multiple objects, granting no - * read leases + * @brief Probe multiple objects, optionally leasing the last complete + * candidate. LastHitOnly leases the last all-true candidate. * @param keys Vector of keys to check - * @return Vector of existence results for each key + * @return Per-key existence results in input order */ std::vector> BatchProbeKey( - const std::vector& keys); + const std::vector& keys, + const GrantLeasePolicy& policy = {}); /** * @brief Create a copy task to copy an object's replicas to target segments diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index 5a4777a030..7b922b55bd 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -185,7 +185,8 @@ class DummyClient : public PyClient { int probeKey(const std::string &key); - std::vector batchProbeKey(const std::vector &keys); + std::vector batchProbeKey(const std::vector &keys, + const GrantLeasePolicy &policy = {}); int64_t getSize(const std::string &key); diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 008dfd9f89..aee43a476c 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -152,13 +152,14 @@ class MasterClient { const std::vector& object_keys); /** - * @brief Point-in-time existence check for multiple objects, granting no - * read leases + * @brief Probe multiple objects, optionally leasing the last complete + * candidate. LastHitOnly leases the last all-true candidate in the results. * @param object_keys Vector of keys to query - * @return Vector containing existence status for each key + * @return Per-key existence results in input order */ [[nodiscard]] std::vector> BatchProbeKey( - const std::vector& object_keys); + const std::vector& object_keys, + const GrantLeasePolicy& policy = {}); /** * @brief Calculate Store-observed cache reuse metrics diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 93afd41c35..c80f92b68c 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -262,7 +262,8 @@ class MasterService { -> tl::expected; std::vector> BatchProbeKey( - const std::vector& keys, const TenantId& tenant_id); + const GrantLeasePolicy& policy, const std::vector& keys, + const TenantId& tenant_id); /** * @brief Fetch all keys for a single tenant. diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h index 6fe9d3dff8..283dd6ef6a 100644 --- a/mooncake-store/include/pyclient.h +++ b/mooncake-store/include/pyclient.h @@ -538,7 +538,8 @@ class PyClient { virtual int probeKey(const std::string &key) = 0; virtual std::vector batchProbeKey( - const std::vector &keys) = 0; + const std::vector &keys, + const GrantLeasePolicy &policy = {}) = 0; virtual int64_t getSize(const std::string &key) = 0; diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 6e5e2ff8e9..19bbc9d7fc 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -387,13 +387,14 @@ class RealClient : public PyClient { int probeKey(const std::string &key); /** - * @brief Point-in-time existence check for multiple objects, granting - * no read leases + * @brief Probe multiple objects, optionally leasing the last complete + * candidate. * @param keys Vector of keys to check - * @return Vector of existence results: 1 if exists, 0 if not exists, -1 - * if error + * @return Per-key existence results. LastHitOnly leases the last all-1 + * candidate. Negative values indicate errors. */ - std::vector batchProbeKey(const std::vector &keys); + std::vector batchProbeKey(const std::vector &keys, + const GrantLeasePolicy &policy = {}); /** * @brief Get the size of an object @@ -801,7 +802,7 @@ class RealClient : public PyClient { tl::expected probeKey_internal(const std::string &key); std::vector> batchProbeKey_internal( - const std::vector &keys); + const std::vector &keys, const GrantLeasePolicy &policy); tl::expected getSize_internal(const std::string &key); diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index c334157923..ba326e63ed 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -54,7 +54,7 @@ class WrappedMasterService { const std::string& tenant_id = "default"); std::vector> BatchProbeKey( - const std::vector& keys, + const GrantLeasePolicy& policy, const std::vector& keys, const std::string& tenant_id = "default"); tl::expected< diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index 0536bd7cab..d26a6126f6 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -20,6 +20,15 @@ namespace mooncake { +enum class ProbeLeaseMode : uint8_t { None = 0, LastHitOnly = 1 }; + +// Candidates are consecutive groups of candidate_size keys in caller order. +// None preserves the lease-free existence probe and ignores candidate_size. +struct GrantLeasePolicy { + ProbeLeaseMode lease_mode{ProbeLeaseMode::None}; + uint64_t candidate_size{1}; +}; + // Constants static constexpr uint64_t WRONG_VERSION = 0; static constexpr uint64_t DEFAULT_VALUE = UINT64_MAX; diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 6a8be15706..b6f8a214d2 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -4125,8 +4125,8 @@ std::vector> Client::BatchIsExist( } std::vector> Client::BatchProbeKey( - const std::vector& keys) { - auto response = master_client_.BatchProbeKey(keys); + const std::vector& keys, const GrantLeasePolicy& policy) { + auto response = master_client_.BatchProbeKey(keys, policy); // Check if we got the expected number of responses if (response.size() != keys.size()) { diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index e17d4e4901..3667ee3cf5 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -1333,10 +1333,10 @@ int DummyClient::probeKey(const std::string& key) { } std::vector DummyClient::batchProbeKey( - const std::vector& keys) { + const std::vector& keys, const GrantLeasePolicy& policy) { auto internal_results = - invoke_batch_rpc<&RealClient::batchProbeKey_internal, bool>(keys.size(), - keys); + invoke_batch_rpc<&RealClient::batchProbeKey_internal, bool>( + keys.size(), keys, policy); std::vector results; results.reserve(internal_results.size()); diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index bf8a1a92dd..4bdcd2181b 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -565,12 +565,13 @@ tl::expected MasterClient::ProbeKey( } std::vector> MasterClient::BatchProbeKey( - const std::vector& object_keys) { + const std::vector& object_keys, + const GrantLeasePolicy& policy) { ScopedVLogTimer timer(1, "MasterClient::BatchProbeKey"); timer.LogRequest("keys_count=", object_keys.size()); auto result = invoke_batch_rpc<&WrappedMasterService::BatchProbeKey, bool>( - object_keys.size(), object_keys, tenant_id_.value()); + object_keys.size(), policy, object_keys, tenant_id_.value()); timer.LogResponse("result=", result.size(), " keys"); return result; } diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 2d10751fa5..c89bdba2b4 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -3023,8 +3023,55 @@ std::vector> MasterService::BatchExistKey( } std::vector> MasterService::BatchProbeKey( - const std::vector& keys, const TenantId& tenant_id) { - return BatchExistKeyImpl(keys, tenant_id, /*grant_lease=*/false); + const GrantLeasePolicy& policy, const std::vector& keys, + const TenantId& tenant_id) { + if (keys.empty()) return {}; + std::string_view invalid_reason; + if (policy.lease_mode != ProbeLeaseMode::None && + policy.lease_mode != ProbeLeaseMode::LastHitOnly) { + invalid_reason = "unknown lease mode"; + } else if (policy.lease_mode == ProbeLeaseMode::LastHitOnly) { + if (policy.candidate_size == 0) { + invalid_reason = "candidate_size is zero"; + } else if (keys.size() % policy.candidate_size != 0) { + invalid_reason = "key count is not divisible by candidate_size"; + } + } + if (!invalid_reason.empty()) { + LOG(WARNING) << "BatchProbeKey: " << invalid_reason + << ", lease_mode=" << static_cast(policy.lease_mode) + << ", candidate_size=" << policy.candidate_size + << ", key_count=" << keys.size(); + return std::vector>( + keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } + + auto results = BatchExistKeyImpl(keys, tenant_id, /*grant_lease=*/false); + if (policy.lease_mode == ProbeLeaseMode::None) return results; + + const auto is_readable = [](const auto& result) { + return result.has_value() && result.value(); + }; + for (size_t end = keys.size(); end > 0; end -= policy.candidate_size) { + const size_t begin = end - policy.candidate_size; + if (!std::all_of(results.begin() + begin, results.begin() + end, + is_readable)) { + continue; + } + + const std::vector candidate(keys.begin() + begin, + keys.begin() + end); + auto leased = + BatchExistKeyImpl(candidate, tenant_id, /*grant_lease=*/true); + // Recheck under the existing per-shard locks. Publish the fresh + // results so a candidate lost since probing cannot remain all-true. + // Failed candidates may retain partial leases until their TTL expires. + std::copy(leased.begin(), leased.end(), results.begin() + begin); + if (std::all_of(leased.begin(), leased.end(), is_readable)) { + break; + } + } + return results; } std::vector> MasterService::BatchExistKeyImpl( diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index b8ecba6204..94ae09d44a 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -2884,9 +2884,9 @@ int RealClient::probeKey(const std::string &key) { } } -std::vector RealClient::batchProbeKey( - const std::vector &keys) { - auto internal_results = batchProbeKey_internal(keys); +std::vector RealClient::batchProbeKey(const std::vector &keys, + const GrantLeasePolicy &policy) { + auto internal_results = batchProbeKey_internal(keys, policy); std::vector results; results.reserve(internal_results.size()); @@ -6725,7 +6725,7 @@ tl::expected RealClient::probeKey_internal( } std::vector> RealClient::batchProbeKey_internal( - const std::vector &keys) { + const std::vector &keys, const GrantLeasePolicy &policy) { if (!client_) { LOG(ERROR) << "Client is not initialized"; return std::vector>( @@ -6738,7 +6738,7 @@ std::vector> RealClient::batchProbeKey_internal( } // Call client BatchProbeKey and return the vector directly - return client_->BatchProbeKey(keys); + return client_->BatchProbeKey(keys, policy); } int RealClient::put_from_with_metadata(const std::string &key, void *buffer, diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index a702b709bf..de169b404e 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -181,7 +181,8 @@ tl::expected WrappedMasterService::ProbeKey( } std::vector> WrappedMasterService::BatchProbeKey( - const std::vector& keys, const std::string& tenant_id) { + const GrantLeasePolicy& policy, const std::vector& keys, + const std::string& tenant_id) { ScopedVLogTimer timer(1, "BatchProbeKey"); timer.LogRequest("keys_count=", keys.size()); @@ -189,7 +190,8 @@ std::vector> WrappedMasterService::BatchProbeKey( master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) : TenantId::kDefaultValue, keys.size(), [&](const TenantId& resolved_tenant_id) { - return master_service_.BatchProbeKey(keys, resolved_tenant_id); + return master_service_.BatchProbeKey(policy, keys, + resolved_tenant_id); }); } diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 5c0d82a325..00fb1a2a7a 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -77,6 +77,8 @@ endif() add_subdirectory(master_service) add_store_test(master_service_test master_service_test.cpp) +add_store_test(master_service_probe_policy_test + master_service_probe_policy_test.cpp) add_store_test(master_metadata_config_test master_metadata_config_test.cpp) add_store_test(dfs_enablement_config_test dfs_enablement_config_test.cpp) add_store_test(master_service_weight_management_test diff --git a/mooncake-store/tests/dummy_client_probe_key_test.cpp b/mooncake-store/tests/dummy_client_probe_key_test.cpp index b8ba20c60d..96dfe99de2 100644 --- a/mooncake-store/tests/dummy_client_probe_key_test.cpp +++ b/mooncake-store/tests/dummy_client_probe_key_test.cpp @@ -150,6 +150,25 @@ TEST_F(DummyClientProbeKeyTest, BatchProbeKeyOverRpcReportsExistence) { EXPECT_EQ(results[2], 1) << key_b << " should exist"; } +TEST_F(DummyClientProbeKeyTest, LastHitOnlyPolicySurvivesBothRpcBoundaries) { + ASSERT_TRUE(SetupStack()); + for (const auto &key : {"a", "b", "c"}) PutData(key, "state"); + const GrantLeasePolicy policy{ProbeLeaseMode::LastHitOnly, 2}; + EXPECT_EQ(dummy_client_->batchProbeKey({"a", "b", "c", "missing"}, policy), + (std::vector{1, 1, 1, 0})); + EXPECT_EQ(dummy_client_->batchProbeKey({"a", "missing", "b", "c"}, policy), + (std::vector{1, 0, 1, 1})); + EXPECT_EQ(real_client_->batchProbeKey({"a", "b", "b", "c"}, policy), + (std::vector{1, 1, 1, 1})); + EXPECT_EQ(dummy_client_->batchProbeKey({"a", "b", "c"}), + (std::vector{1, 1, 1})); + EXPECT_TRUE(dummy_client_->batchProbeKey({}, policy).empty()); + const auto invalid = dummy_client_->batchProbeKey( + {"a", "b"}, {ProbeLeaseMode::LastHitOnly, 0}); + ASSERT_EQ(invalid.size(), 2u); + for (int value : invalid) EXPECT_LT(value, 0); +} + } // namespace testing } // namespace mooncake diff --git a/mooncake-store/tests/master_service_probe_policy_test.cpp b/mooncake-store/tests/master_service_probe_policy_test.cpp new file mode 100644 index 0000000000..bcb88567b5 --- /dev/null +++ b/mooncake-store/tests/master_service_probe_policy_test.cpp @@ -0,0 +1,284 @@ +#include "master_service_test_fixture.h" +#include "rpc_service.h" + +#include +#include +#include + +namespace mooncake::test { + +class BatchProbePolicyTest : public MasterServiceTest { + protected: + static constexpr GrantLeasePolicy kLastPair{ProbeLeaseMode::LastHitOnly, 2}; + + std::chrono::system_clock::time_point Deadline(MasterService& service, + const std::string& key) { + MasterServiceTestPeer::MetadataAccessorRO accessor( + &service, + MasterServiceTestPeer(service).MakeObjectIdentityForRequest( + key, TenantId::Default())); + EXPECT_TRUE(accessor.Exists()); + return accessor.Exists() ? accessor.Get().EvictionDeadline() + : std::chrono::system_clock::time_point{}; + } + + void ExpectResults(const std::vector>& result, + const std::vector& expected) { + ASSERT_EQ(result.size(), expected.size()); + for (size_t i = 0; i < result.size(); ++i) { + ASSERT_TRUE(result[i].has_value()) << i; + EXPECT_EQ(result[i].value(), expected[i]) << i; + } + } +}; + +TEST_F(BatchProbePolicyTest, NonePreservesReadAndSoftPinDeadlines) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(60000).build()); + const auto context = PrepareSimpleSegment(service); + ReplicateConfig config; + config.replica_num = 1; + config.soft_pin_action = SoftPinAction::ENABLE; + PutCompletedObject(service, context.client_id, "key", config); + const auto before = Deadline(service, "key"); + const auto soft_pin = GetSoftPinDeadline(service, "key"); + ASSERT_TRUE(soft_pin.has_value()); + + // candidate_size is deliberately ignored in None mode. + ExpectResults( + service.BatchProbeKey({ProbeLeaseMode::None, 0}, + {"key", "missing", "key"}, TenantId::Default()), + {true, false, true}); + EXPECT_EQ(Deadline(service, "key"), before); + EXPECT_EQ(GetSoftPinDeadline(service, "key"), soft_pin); + + ExpectResults(service.BatchExistKey({"key"}, TenantId::Default()), {true}); + EXPECT_GT(Deadline(service, "key"), before); + const auto leased = Deadline(service, "key"); + ExpectResults(service.BatchProbeKey({}, {"key"}, TenantId::Default()), + {true}); + EXPECT_EQ(Deadline(service, "key"), leased); +} + +TEST_F(BatchProbePolicyTest, SelectsLastCompleteSparseCandidate) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(60000).build()); + const auto context = PrepareSimpleSegment(service); + ReplicateConfig config; + config.replica_num = 1; + config.soft_pin_action = SoftPinAction::ENABLE; + const std::vector keys = {"p0-state", "p0-conv", "p1-state", + "p1-conv", "p2-state", "p2-conv", + "p3-state", "p3-conv"}; + for (size_t i : {0u, 1u, 4u, 5u, 6u}) { + PutCompletedObject(service, context.client_id, keys[i], config); + } + const auto old = Deadline(service, keys[0]); + const auto partial = Deadline(service, keys[6]); + const auto chosen = Deadline(service, keys[4]); + const auto soft_pin = GetSoftPinDeadline(service, keys[4]); + ExpectResults(service.BatchProbeKey(kLastPair, keys, TenantId::Default()), + {true, true, false, false, true, true, true, false}); + EXPECT_EQ(Deadline(service, keys[0]), old); + EXPECT_EQ(Deadline(service, keys[6]), partial); + EXPECT_GT(Deadline(service, keys[4]), chosen); + EXPECT_EQ(GetSoftPinDeadline(service, keys[4]), soft_pin); + for (size_t i : {4u, 5u}) { + auto removed = service.Remove(keys[i], TenantId::Default()); + ASSERT_FALSE(removed.has_value()); + EXPECT_EQ(removed.error(), ErrorCode::OBJECT_HAS_LEASE); + } +} + +TEST_F(BatchProbePolicyTest, OnlySelectedCandidateSurvivesNormalEviction) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(60000).build()); + const auto context = PrepareSimpleSegment(service); + ReplicateConfig config; + config.replica_num = 1; + const std::vector keys = {"z", "y", "b", "a"}; + for (const auto& key : keys) { + PutCompletedObject(service, context.client_id, key, config); + } + ExpectResults(service.BatchProbeKey(kLastPair, keys, TenantId::Default()), + {true, true, true, true}); + MasterServiceTestPeer(service).RunBatchEvictForTesting(1.0, 1.0); + ExpectResults(service.BatchProbeKey({}, keys, TenantId::Default()), + {false, false, true, true}); +} + +TEST_F(BatchProbePolicyTest, IncompleteAndUnreadableCandidatesAreNotLeased) { + MasterService service; + const auto context = PrepareSimpleSegment(service); + ReplicateConfig config; + config.replica_num = 1; + PutCompletedObject(service, context.client_id, "complete", config); + ASSERT_TRUE(service + .PutStart(context.client_id, "writing", TenantId::Default(), + 1024, config) + .has_value()); + const auto before = Deadline(service, "complete"); + const std::vector keys = {"complete", "writing"}; + ExpectResults(service.BatchProbeKey(kLastPair, keys, TenantId::Default()), + {true, false}); + EXPECT_EQ(Deadline(service, "complete"), before); + ASSERT_TRUE(service + .PutEnd(context.client_id, "writing", TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + ASSERT_TRUE(service.UnmountSegment(context.segment_id, context.client_id) + .has_value()); + ExpectResults(service.BatchProbeKey(kLastPair, keys, TenantId::Default()), + {false, false}); +} + +TEST_F(BatchProbePolicyTest, ValidationDoesNotLeaseAndEmptyInputIsEmpty) { + MasterService service; + const auto context = PrepareSimpleSegment(service); + PutCompletedObject(service, context.client_id, "key", {.replica_num = 1}); + const auto before = Deadline(service, "key"); + for (const auto& policy : + std::vector{{ProbeLeaseMode::LastHitOnly, 0}, + {ProbeLeaseMode::LastHitOnly, 3}, + {ProbeLeaseMode::LastHitOnly, + std::numeric_limits::max()}, + {static_cast(99), 1}}) { + auto result = + service.BatchProbeKey(policy, {"key", "key"}, TenantId::Default()); + ASSERT_EQ(result.size(), 2u); + for (const auto& item : result) { + ASSERT_FALSE(item.has_value()); + EXPECT_EQ(item.error(), ErrorCode::INVALID_PARAMS); + } + EXPECT_EQ(Deadline(service, "key"), before); + } + EXPECT_TRUE( + service.BatchProbeKey(kLastPair, {}, TenantId::Default()).empty()); + EXPECT_TRUE(service.BatchProbeKey({}, {}, TenantId::Default()).empty()); +} + +TEST_F(BatchProbePolicyTest, DuplicateKeysAndSameShardPreserveExistence) { + MasterService service; + const auto context = PrepareSimpleSegment(service); + const std::string first = "same-shard-first"; + std::string second; + for (int i = 0; i < 100000; ++i) { + const std::string candidate = "same-shard-" + std::to_string(i); + if (MetadataShardIndex(service, candidate) == + MetadataShardIndex(service, first)) { + second = candidate; + break; + } + } + ASSERT_FALSE(second.empty()); + for (const auto& key : {first, second}) { + PutCompletedObject(service, context.client_id, key, {.replica_num = 1}); + } + ExpectResults( + service.BatchProbeKey(kLastPair, {first, first, first, second}, + TenantId::Default()), + {true, true, true, true}); + ExpectResults( + service.BatchProbeKey(kLastPair, {"missing", first, first, first}, + TenantId::Default()), + {false, true, true, true}); + ExpectResults( + service.BatchProbeKey({ProbeLeaseMode::LastHitOnly, 1}, + {first, "missing", second}, TenantId::Default()), + {true, false, true}); +} + +TEST_F(BatchProbePolicyTest, PreservesExistingSharedGroupLease) { + MasterService service; + const auto context = PrepareSimpleSegment(service); + ReplicateConfig config; + config.replica_num = 1; + config.group_ids = std::vector{"shared-group"}; + for (const auto& key : {"old", "selected"}) { + PutCompletedObject(service, context.client_id, key, config); + } + const auto before = Deadline(service, "old"); + ExpectResults( + service.BatchProbeKey({}, {"old", "selected"}, TenantId::Default()), + {true, true}); + EXPECT_EQ(Deadline(service, "old"), before); + ExpectResults( + service.BatchProbeKey({ProbeLeaseMode::LastHitOnly, 1}, + {"old", "selected"}, TenantId::Default()), + {true, true}); + EXPECT_GT(Deadline(service, "old"), before); + EXPECT_EQ(Deadline(service, "old"), Deadline(service, "selected")); +} + +TEST_F(BatchProbePolicyTest, TenantIsolationAndErrorsThroughWrapper) { + WrappedMasterService service(MakeStrictWrappedConfig({"left", "right"})); + const auto segment = MakeSegment("tenant-probe-segment"); + const UUID client_id = generate_uuid(); + ASSERT_TRUE(service.MountSegment(segment, client_id).has_value()); + const std::vector keys = {"a", "b", "c", "d"}; + for (const auto& tenant : {std::string("left"), std::string("right")}) { + const size_t offset = tenant == "left" ? 0 : 2; + const std::vector present = {keys[offset], + keys[offset + 1]}; + auto started = service.BatchPutStart(client_id, present, {1024, 1024}, + {.replica_num = 1}, tenant); + for (const auto& item : started) ASSERT_TRUE(item.has_value()); + auto ended = service.BatchPutEnd(client_id, MakeObjectMetas(present), + ReplicaType::MEMORY, tenant); + for (const auto& item : ended) ASSERT_TRUE(item.has_value()); + } + ExpectResults(service.BatchProbeKey(kLastPair, keys, "left"), + {true, true, false, false}); + ExpectResults(service.BatchProbeKey(kLastPair, keys, "right"), + {false, false, true, true}); + ExpectResults(service.BatchProbeKey(kLastPair, keys, "unregistered"), + {false, false, false, false}); + auto invalid = service.BatchProbeKey(kLastPair, keys, "_invalid-tenant"); + ASSERT_EQ(invalid.size(), keys.size()); + for (const auto& item : invalid) { + ASSERT_FALSE(item.has_value()); + EXPECT_EQ(item.error(), ErrorCode::INVALID_PARAMS); + } +} + +TEST_F(BatchProbePolicyTest, + ConcurrentRemovalPreservesLeasedSelectionOrFallback) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(60000).build()); + const auto context = PrepareSimpleSegment(service); + for (int iteration = 0; iteration < 64; ++iteration) { + std::vector keys; + for (int part = 0; part < 8; ++part) { + keys.push_back("race-" + std::to_string(iteration) + "-" + + std::to_string(part)); + PutCompletedObject(service, context.client_id, keys.back(), + {.replica_num = 1}); + } + std::barrier start(2); + std::thread remover([&] { + start.arrive_and_wait(); + for (size_t i = 4; i < keys.size(); ++i) + service.Remove(keys[i], TenantId::Default()); + }); + start.arrive_and_wait(); + auto result = service.BatchProbeKey({ProbeLeaseMode::LastHitOnly, 4}, + keys, TenantId::Default()); + remover.join(); + ASSERT_EQ(result.size(), keys.size()); + for (const auto& item : result) ASSERT_TRUE(item.has_value()); + for (size_t i = 0; i < 4; ++i) EXPECT_TRUE(result[i].value()); + const bool last_complete = + std::all_of(result.begin() + 4, result.end(), + [](const auto& item) { return item.value(); }); + const size_t selected = last_complete ? 4 : 0; + // Partial results and partial leases on the failed candidate are + // allowed. The last all-true candidate must nevertheless be protected. + for (size_t i = selected; i < selected + 4; ++i) { + auto removed = service.Remove(keys[i], TenantId::Default()); + ASSERT_FALSE(removed.has_value()); + EXPECT_EQ(removed.error(), ErrorCode::OBJECT_HAS_LEASE); + } + } +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index ccccfe7180..04f2387804 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -2675,7 +2675,7 @@ TEST_F(MasterServiceTest, BatchProbeKeyReportsPointInTimeExistence) { const std::string missing_key = "probe_batch_missing_key"; auto results = service_->BatchProbeKey( - {existing_key, missing_key, existing_key}, TenantId::Default()); + {}, {existing_key, missing_key, existing_key}, TenantId::Default()); ASSERT_EQ(3u, results.size()); ASSERT_TRUE(results[0].has_value()); EXPECT_TRUE(*results[0]); diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp index 75cdfdbe2c..2061c1ffa9 100644 --- a/mooncake-store/tests/pybind_client_test.cpp +++ b/mooncake-store/tests/pybind_client_test.cpp @@ -340,6 +340,90 @@ TEST_F(RealClientTest, SessionRangesReadDfsAndPropagateShortRead) { } #ifdef MOONCAKE_TEST_CUDA_H2D +TEST_F(RealClientTest, LastHitOnlyGpuCheckpointRoundTrip) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device is unavailable"; + } + + ScopedEnvVar local_memcpy("MC_STORE_MEMCPY", "1"); + StartMasterAndSetupClient(); + ASSERT_FALSE(HasFatalFailure()); + + // Three synthetic TP8 checkpoints, each with two state components per + // rank. The last checkpoint is missing one component. + constexpr size_t kCandidateSize = 16; + constexpr size_t kKeyCount = 3 * kCandidateSize; + constexpr size_t kPartSize = 4096; + constexpr size_t kPayloadSize = kKeyCount * kPartSize; + std::vector source(kPayloadSize); + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i % 251); + } + void* allocation = nullptr; + ASSERT_EQ(cudaMalloc(&allocation, 2 * kPayloadSize), cudaSuccess); + bool registered = false; + auto cleanup = [this, ®istered](void* ptr) { + if (registered) EXPECT_EQ(py_client_->unregister_buffer(ptr), 0); + EXPECT_EQ(cudaFree(ptr), cudaSuccess); + }; + std::unique_ptr gpu_owner(allocation, cleanup); + ASSERT_EQ(cudaMemcpy(allocation, source.data(), kPayloadSize, + cudaMemcpyHostToDevice), + cudaSuccess); + ASSERT_EQ(py_client_->register_buffer(allocation, 2 * kPayloadSize), 0); + registered = true; + + std::vector keys; + std::vector buffers; + for (size_t i = 0; i < kKeyCount; ++i) { + keys.push_back("gpu-probe-" + std::to_string(i)); + if (i + 1 < kKeyCount) { + buffers.push_back(static_cast(allocation) + i * kPartSize); + } + } + const std::vector present(keys.begin(), keys.end() - 1); + ASSERT_EQ( + py_client_->batch_put_from( + present, buffers, std::vector(present.size(), kPartSize), + {.replica_num = 1}), + std::vector(present.size(), 0)); + + auto expected = std::vector(kKeyCount, 1); + expected.back() = 0; + EXPECT_EQ(py_client_->batchProbeKey(keys), expected); + ASSERT_EQ(py_client_->batchProbeKey( + keys, {ProbeLeaseMode::LastHitOnly, kCandidateSize}), + expected); + + const std::vector selected(keys.begin() + kCandidateSize, + keys.begin() + 2 * kCandidateSize); + for (const auto& key : selected) { + EXPECT_EQ(py_client_->remove(key), toInt(ErrorCode::OBJECT_HAS_LEASE)); + } + // Existence results remain true for unselected keys without leasing them. + EXPECT_EQ(py_client_->remove(keys.front()), 0); + EXPECT_EQ(py_client_->remove(keys[kKeyCount - 2]), 0); + + auto* destination = static_cast(allocation) + kPayloadSize; + ASSERT_EQ(cudaMemset(destination, 0, kCandidateSize * kPartSize), + cudaSuccess); + buffers.clear(); + for (size_t i = 0; i < kCandidateSize; ++i) { + buffers.push_back(destination + i * kPartSize); + } + ASSERT_EQ( + py_client_->batch_get_into( + selected, buffers, std::vector(kCandidateSize, kPartSize)), + std::vector(kCandidateSize, kPartSize)); + std::vector actual(kCandidateSize * kPartSize); + ASSERT_EQ(cudaMemcpy(actual.data(), destination, actual.size(), + cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_TRUE(std::equal(actual.begin(), actual.end(), + source.begin() + kCandidateSize * kPartSize)); +} + TEST_F(RealClientTest, SessionRangesReadDfsIntoGpuUsesPinnedArena) { int device_count = 0; if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { diff --git a/mooncake-wheel/tests/test_batch_probe_policy.py b/mooncake-wheel/tests/test_batch_probe_policy.py new file mode 100644 index 0000000000..96e13cb701 --- /dev/null +++ b/mooncake-wheel/tests/test_batch_probe_policy.py @@ -0,0 +1,85 @@ +"""CPU/TCP policy checks against a real Master, optionally via the client daemon. + +Set MASTER_SERVER to an isolated Master address. Set REAL_CLIENT_ADDRESS to +also run this suite through the production DummyClient/RealClient RPC path. +""" + +import os +import unittest + +from mooncake.store import MooncakeDistributedStore + + +class TestBatchProbePolicy(unittest.TestCase): + def setUp(self): + self.store = MooncakeDistributedStore() + self.addCleanup(self.store.close) + daemon = os.getenv("REAL_CLIENT_ADDRESS") + if daemon: + status = self.store.setup_dummy(16 << 20, 16 << 20, daemon) + else: + status = self.store.setup( + "127.0.0.1:0", + "P2PHANDSHAKE", + 32 << 20, + 16 << 20, + "tcp", + "", + os.getenv("MASTER_SERVER", "127.0.0.1:50051"), + ) + self.assertEqual(status, 0) + self.keys = [f"probe-policy-{os.getpid()}-{self.id()}-{i}" for i in range(4)] + for key in self.keys[:3]: + self.assertEqual(self.store.put(key, b"checkpoint"), 0) + self.addCleanup(self.store.remove, key, True) + + def test_existence_and_last_complete_candidate(self): + a, b, c, missing = self.keys + self.assertEqual(self.store.batch_probe_key(self.keys), [1, 1, 1, 0]) + self.assertEqual(self.store.batch_is_exist(self.keys), [1, 1, 1, 0]) + self.assertEqual( + self.store.batch_probe_key(self.keys, "LastHitOnly", 2), [1, 1, 1, 0] + ) + self.assertEqual( + self.store.batch_probe_key([a, b, b, c], "LastHitOnly", 2), + [1, 1, 1, 1], + ) + self.assertEqual( + self.store.batch_probe_key([a, missing, b, c], "LastHitOnly", 2), + [1, 0, 1, 1], + ) + self.assertEqual( + self.store.batch_probe_key([a, missing], "LastHitOnly", 2), [1, 0] + ) + self.assertEqual(self.store.get(c), b"checkpoint") + + def test_validation_duplicates_and_empty_input(self): + a, b, _, _ = self.keys + self.assertEqual( + self.store.batch_probe_key([a, b, a, a], "LastHitOnly", 2), + [1, 1, 1, 1], + ) + self.assertEqual(self.store.batch_probe_key([a, b], "LastHitOnly"), [1, 1]) + self.assertEqual(self.store.batch_probe_key([a], candidate_size=0), [1]) + self.assertEqual(self.store.batch_probe_key([]), []) + self.assertEqual(self.store.batch_probe_key([], "LastHitOnly", 0), []) + for size in (0, 3): + result = self.store.batch_probe_key([a, b], "LastHitOnly", size) + self.assertEqual(len(result), 2) + self.assertTrue(all(value < 0 for value in result)) + with self.assertRaises(ValueError): + self.store.batch_probe_key([a], "unknown") + with self.assertRaises(TypeError): + self.store.batch_probe_key([a], "LastHitOnly", -1) + + def test_uninitialized_errors_are_not_misses(self): + uninitialized = MooncakeDistributedStore() + self.addCleanup(uninitialized.close) + for policy in ("none", "LastHitOnly"): + result = uninitialized.batch_probe_key(self.keys, policy, 2) + self.assertEqual(len(result), len(self.keys)) + self.assertTrue(all(value < 0 for value in result)) + + +if __name__ == "__main__": + unittest.main() From d750cddafa4b62387ca135949595008c51e05815 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Sat, 17 Oct 2026 01:46:00 +0800 Subject: [PATCH 02/14] [Doc] Refresh ecosystem integrations and archive older README news (#4659) Signed-off-by: Teng Ma --- README.md | 109 +++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 83 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index e0da1235a3..5623408399 100644 --- a/README.md +++ b/README.md @@ -33,46 +33,61 @@ Under real workloads, Mooncake’s innovative architecture enables Kimi to handl

πŸ”„ Updates

+ + +- **Oct 10, 2026**: TensorRT-LLM merged a native [Mooncake Store connector](https://github.com/NVIDIA/TensorRT-LLM/pull/19171) into KVCMv2, enabling content-addressed KV cache pooling and prefix reuse across serving instances. +- **Oct 6, 2026**: [TokenSpeed v0.1.1](https://github.com/lightseekorg/tokenspeed/releases/tag/v0.1.1) includes Mooncake Store as an L3 KV cache backend beneath its host cache, alongside Mooncake Transfer Engine support for PD and EPD disaggregation. +- **Sept 17, 2026**: vLLM added [cross-encoder embedding reuse](https://github.com/vllm-project/vllm/pull/56242) through Mooncake Store, complementing its [encoder-cache transfer connector](https://github.com/vllm-project/vllm/pull/41567) for multimodal disaggregated serving. Both capabilities are included in [vLLM v0.31.0](https://github.com/vllm-project/vllm/releases/tag/v0.31.0). +- **Aug 28, 2026**: LMDeploy merged [Mooncake Store support](https://github.com/InternLM/lmdeploy/pull/4903), enabling shared KV cache pooling and prefix reuse across PyTorch inference instances. Included in [LMDeploy v0.18.0](https://github.com/InternLM/lmdeploy/releases/tag/v0.18.0). - **Aug 24, 2026**: [AgentX / InferenceX v3](https://newsletter.semianalysis.com/p/agentx-inferencexv3-does-cuda-moat) highlights Mooncake's contributions to agentic inference workloads, including ROCm wheels, CI, and packaging improvements. [Mooncake optimization details](https://inferencex.semianalysis.com/agentx/optimizations/mooncake). -- **Aug 20, 2026**: Mooncake is integrated into [Miles](https://github.com/radixark/miles) as a rollout data-transfer backend for the fragmented, heterogeneous data moving between rollout and training in disaggregated RL. Blogs: [KVCache.AI](https://kvcache.ai/blog/mooncake-rl-rollout-data-transfer/), [lmsys](https://www.lmsys.org/blog/2026-08-20-miles-mooncake-rollout-data-transfer). -- **Aug 17, 2026**: Mooncake is integrated into [Speculators](https://github.com/vllm-project/speculators) as a distributed backend for multi-node online training, efficiently moves hidden-state between vLLM inference workers and Speculators trainers through RDMA, eliminating the need for massive hidden-state storage in offline training. [Benchmark on GB300 NVL72](https://x.com/mgoin_/status/2072785822231728363). -- **May 7, 2026**: πŸš€ [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) β€” a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing! -- **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s β†’ 7.2s) with zero-copy RDMA transfer across thousands of GPUs.
More +- **Aug 20, 2026**: Mooncake is integrated into [Miles](https://github.com/radixark/miles) as a rollout data-transfer backend for the fragmented, heterogeneous data moving between rollout and training in disaggregated RL. Blogs: [KVCache.AI](https://kvcache.ai/blog/mooncake-rl-rollout-data-transfer/), [lmsys](https://www.lmsys.org/blog/2026-08-20-miles-mooncake-rollout-data-transfer). +- **Aug 17, 2026**: Mooncake is integrated into [Speculators](https://github.com/vllm-project/speculators) as a distributed backend for multi-node online training, efficiently moves hidden-state between vLLM inference workers and Speculators trainers through RDMA, eliminating the need for massive hidden-state storage in offline training. [Benchmark on GB300 NVL72](https://x.com/mgoin_/status/2072785822231728363). +- **Aug 5, 2026**: llm-d added a [Mooncake PD deployment guide](https://github.com/llm-d/llm-d/pull/1511), complementing its [Mooncake Store tiered offloading integration](https://github.com/llm-d/llm-d/pull/1826) for vLLM serving on Kubernetes. +- **Aug 1, 2026**: Dynamo merged an [event-driven Mooncake shared-cache index](https://github.com/ai-dynamo/dynamo/pull/11239) for SGLang HiCache routing, using Mooncake object events to track reusable KV pages in the shared pool. +- **Jul 24, 2026**: FlexKV added [Mooncake Store as a distributed REMOTE-tier backend](https://github.com/taco-project/FlexKV/pull/231), enabling content-addressed KV cache sharing alongside its existing Mooncake Transfer Engine integration. +- **Jul 12, 2026**: TransferQueue added [GPUDirect RDMA support for its Mooncake Store backend](https://github.com/Ascend/TransferQueue/pull/131), enabling GPU clients to exchange RL tensors with the host-memory store through GDR. +- **May 7, 2026**: πŸš€ [vLLM officially features Mooncake Store](https://vllm.ai/blog/mooncake-store) β€” a deep dive into how Mooncake's distributed KVCache engine supercharges vLLM inference with high-throughput, memory-efficient, cross-instance KV cache sharing! +- **Apr 29, 2026**: SGLang introduces [RDMA-based P2P weight transfer for large-scale distributed RL](https://lmsys.org/blog/2026-04-29-p2p-update/) using Mooncake TransferEngine, achieving 7x faster weight updates for the 1T-parameter Kimi-K2 model (53s β†’ 7.2s) with zero-copy RDMA transfer across thousands of GPUs. +- **Apr 3, 2026**: LMCache introduced a native [Mooncake Store L2 adapter for multiprocess mode](https://github.com/LMCache/LMCache/pull/2911), connecting the LMCache MP server to the distributed KV cache pool. +- **Apr 1, 2026**: The [TENT technical report](https://arxiv.org/abs/2604.00368) introduces declarative data-movement orchestration for heterogeneous interconnects, with evaluations in SGLang HiCache and checkpoint-engine workloads. - **Mar 19, 2026**: [TorchSpec: Speculative Decoding Training at Scale](https://pytorch.org/blog/torchspec-speculative-decoding-training-at-scale) is [open sourced](https://github.com/torchspec-project/TorchSpec), using Mooncake to decouple inference and training via efficient hidden states management. +- **Mar 16, 2026**: FastDeploy added [Mooncake Store cache read/write support to its PD pipeline](https://github.com/PaddlePaddle/FastDeploy/pull/6768), complementing its [global KV cache pooling](https://github.com/PaddlePaddle/FastDeploy/blob/develop/docs/features/global_cache_pooling.md) across inference instances. - **Mar 5, 2026**: [LightX2V](https://github.com/ModelTC/LightX2V/pull/893) now supports disaggregated deployment based on Mooncake, enabling encoder/transformer service decoupling with Mooncake Transfer Engine for high-performance cross-device and cross-machine data transfer. Details in [blog](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/). +- **Mar 3, 2026**: verl added a [Mooncake checkpoint-engine backend](https://github.com/verl-project/verl/pull/5176) for peer-to-peer weight synchronization between training and rollout workers. - **Feb 25, 2026**: [SGLang](https://github.com/sgl-project/sglang) merged [Encoder Global Cache Manager](https://github.com/sgl-project/sglang/pull/16137), introducing a Mooncake-powered global multimodal embedding cache that enables cross-instance sharing of ViT embeddings to avoid redundant GPU computation. - **Feb 24, 2026**: [vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) introduces disaggregated inference connectors with support for both `MooncakeStoreConnector` and `MooncakeTransferEngineConnector` for multi-node omni-modality pipelines. - **Feb 12, 2026**: [Mooncake Joins PyTorch Ecosystem](https://pytorch.org/blog/mooncake-joins-pytorch-ecosystem/) We are thrilled to announce that Mooncake has officially joined the PyTorch Ecosystem! - **Jan 28, 2026**: [FlexKV](https://github.com/taco-project/FlexKV), a distributed KV store and cache system from Tencent and NVIDIA in collaboration with the community, now supports [distributed KVCache reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md) with the Mooncake Transfer Engine. - - **Dec 27, 2025**: Collaboration with [ROLL](https://github.com/alibaba/ROLL)! Check out the paper [here](https://arxiv.org/abs/2512.22560). - - **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings. - - **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference. - - **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups. - - **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance. - - **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://docs.vllm.ai/projects/ascend/zh-cn/main/user_guide/feature_guide/kv_pool.html). - - **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers. - - **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s. - - **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching. - - **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/disaggregated_prefill.html) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs. - - **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput. - - **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmdeploy.html) for LMDeploy. - - **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md). - - **May 8, 2025**: [Mooncake x LMCache](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmcache/index.html) unite to pioneer KVCache-centric LLM serving system. - - **May 5, 2025**: Supported by Mooncake Team, SGLang release guidance to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs. - - **Apr 22, 2025**: LMCache officially supports Mooncake Store as a remote connector. - - **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer. +- **Dec 27, 2025**: Collaboration with [ROLL](https://github.com/alibaba/ROLL)! Check out the paper [here](https://arxiv.org/abs/2512.22560). +- **Dec 23, 2025**: SGLang introduces [Encode-Prefill-Decode (EPD) Disaggregation](https://lmsys.org/blog/2026-01-12-epd/) with Mooncake as a transfer backend. This integration allows decoupling compute-intensive multimodal encoders (e.g., Vision Transformers) from language model nodes, utilizing Mooncake's RDMA engine for zero-copy transfer of large multimodal embeddings. +- **Dec 19, 2025**: Mooncake Transfer Engine has been [integrated into TensorRT LLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils) for KVCache transfer in PD-disaggregated inference. +- **Dec 19, 2025**: Mooncake Transfer Engine has been directly integrated into vLLM v1 as a [KV Connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/) in PD-disaggregated setups. +- **Nov 07, 2025**: [RBG + SGLang HiCache + Mooncake](https://github.com/sgl-project/rbg/blob/main/keps/74-mooncake-integration/README.md), a role-based out-of-the-box solution for cloud native deployment, which is elastic, scalable, and high-performance. +- **Oct 15, 2025**: SGLang introduced [Mooncake EP and a collective communication backend](https://github.com/sgl-project/sglang/pull/10423) as the foundation for fault-tolerant Elastic Expert Parallel inference. +- **Sept 18, 2025**: Mooncake Store empowers vLLM Ascend by serving as [the distributed KV cache pool backend](https://github.com/vllm-project/vllm-ascend/blob/main/docs/source/user_guide/feature_guide/kv_cache_pool_ascend_store.md). +- **Sept 10, 2025**: SGLang officially supports Mooncake Store as a [hierarchical KV caching storage backend](https://lmsys.org/blog/2025-09-10-sglang-hicache/). The integration extends RadixAttention with multi-tier KV cache storage across device, host, and remote storage layers. +- **Sept 10, 2025**: The official & high-performance version of Mooncake P2P Store is open-sourced as [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine/). It has been successfully applied in K1.5 and K2 production training, updating Kimi-K2 model (1T parameters) across thousands of GPUs in ~20s. +- **Aug 23, 2025**: [xLLM](https://github.com/jd-opensource/xllm) high-performance inference engine builds hybrid KV cache management based on Mooncake, supporting global KV cache management with intelligent offloading and prefetching. +- **Aug 18, 2025**: vLLM-Ascend [integrates Mooncake Transfer Engine](https://github.com/vllm-project/vllm-ascend/blob/main/docs/source/tutorials/features/pd_disaggregation_mooncake_multi_node.md) for KV cache register and disaggregate prefill, enabling efficient distributed inference on Ascend NPUs. +- **Jul 20, 2025**: Mooncake powers [the deployment of Kimi K2](https://lmsys.org/blog/2025-07-20-k2-large-scale-ep/) on 128 H200 GPUs with PD disaggregation and large-scale expert parallelism, achieving 224k tokens/sec prefill throughput and 288k tokens/sec decode throughput. +- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmdeploy.html) for LMDeploy. +- **May 9, 2025**: NIXL officially supports Mooncake Transfer Engine as [a backend plugin](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md). +- **May 8, 2025**: [Mooncake x LMCache](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmcache/index.html) unite to pioneer KVCache-centric LLM serving system. +- **May 5, 2025**: Supported by Mooncake Team, SGLang release guidance to deploy DeepSeek with PD Disaggregation on 96 H100 GPUs. +- **Apr 22, 2025**: LMCache officially supports Mooncake Store as a remote connector. +- **Apr 10, 2025**: SGLang officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer. - **Mar 7, 2025**: We open-sourced the Mooncake Store, a distributed KVCache based on Transfer Engine. vLLM's xPyD disaggregated prefilling & decoding based on Mooncake Store will be released soon. - - **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**! - - **Feb 21, 2025**: The updated traces used in our FAST'25 paper have been released. - - **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer. +- **Feb 25, 2025**: Mooncake receives the **Best Paper Award** at **FAST 2025**! +- **Feb 21, 2025**: The updated traces used in our FAST'25 paper have been released. +- **Dec 16, 2024**: vLLM officially supports Mooncake Transfer Engine for disaggregated prefilling and KV cache transfer. - **Nov 28, 2024**: We open-sourced the Transfer Engine, the central component of Mooncake. We also provide two demonstrations of Transfer Engine: a P2P Store and vLLM integration. - **July 9, 2024**: We open-sourced the trace as a JSONL file. - - **June 27, 2024**: We present a series of Chinese blogs with more discussions on zhihu 1, 2, 3, 4, 5, 6, 7. - - **June 26, 2024**: Initial technical report release. +- **June 27, 2024**: We present a series of Chinese blogs with more discussions on zhihu 1, 2, 3, 4, 5, 6, 7. +- **June 26, 2024**: Initial technical report release.
@@ -193,6 +208,48 @@ Mooncake integrates with [vLLM](https://github.com/vllm-project/vllm) to acceler +

🌐 Ecosystem Integrations

+ +Mooncake connects inference engines, cache middleware, RL systems, and deployment tools. The directory below links to upstream guides and implementations for each integration, checked on Oct 15, 2026. + +
+Integration directory + +**TE** = Transfer Engine; **Store** = Mooncake Store; **EP/PG** = expert-parallel and process-group backends. Capabilities depend on the upstream version and configuration; consult the linked guide or implementation when deploying. + +| Project | Mooncake components | Integration use cases | Guide or implementation | +| --- | --- | --- | --- | +| SGLang | TE, Store, EP/PG | PD/EPD transfer, HiCache, multimodal embedding caching, weight synchronization, and Elastic EP | [Integration guide](https://kvcache-ai.github.io/Mooncake/deployment/integrations/sglang/index.html), [Elastic EP](https://github.com/sgl-project/sglang/pull/10423) | +| vLLM | TE, Store | PD transfer, shared KV pooling, encoder-cache transfer, and cross-encoder embedding reuse | [PD connector](https://docs.vllm.ai/en/latest/features/mooncake_connector_usage/), [Store connector](https://docs.vllm.ai/en/latest/features/mooncake_store_connector_usage/), [encoder connector](https://docs.vllm.ai/en/latest/features/mooncake_ec_connector_usage/) | +| TokenSpeed | TE, Store | PD/EPD transfer and Store-backed L3 beneath the host KV cache | [Server configuration](https://github.com/lightseekorg/tokenspeed/blob/main/docs/configuration/server.md) | +| TensorRT-LLM | TE, Store | PD transfer and shared KV pooling through the KVCMv2 Store connector merged into main | [TE implementation](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/executor/cache_transmission/mooncake_utils), [Store connector](https://github.com/NVIDIA/TensorRT-LLM/pull/19171) | +| LMDeploy | TE, Store | PD transfer and cross-instance KV prefix reuse | [TE guide](https://kvcache-ai.github.io/Mooncake/deployment/integrations/lmdeploy.html), [Store connector](https://github.com/InternLM/lmdeploy/pull/4903) | +| vLLM-Ascend | TE, Store | NPU PD transfer and KV pooling through AscendStoreConnector | [PD deployment](https://github.com/vllm-project/vllm-ascend/blob/main/docs/source/tutorials/features/pd_disaggregation_mooncake_multi_node.md), [KV pool](https://github.com/vllm-project/vllm-ascend/blob/main/docs/source/user_guide/feature_guide/kv_cache_pool_ascend_store.md) | +| FastDeploy | Store | Global KV pooling and cache reuse in PD serving | [Global cache pooling](https://github.com/PaddlePaddle/FastDeploy/blob/develop/docs/features/global_cache_pooling.md) | +| xLLM | TE, Store | KV transfer and distributed cache management | [KV transfer backends](https://github.com/jd-opensource/xllm/tree/main/xllm/core/kv_cache/transfer) | +| Chitu | TE | Multi-prefill/multi-decode KV transfer over RDMA | [PD disaggregation](https://github.com/thu-pacman/chitu/blob/public-main/chitu/distributed/pd_disaggregation/README.md) | +| LMCache | TE, Store | Remote KV storage, a native MP L2 adapter, and P2P cache sharing | [Store guide](https://docs.lmcache.ai/kv_cache/storage_backends/mooncake.html), [MP adapter](https://docs.lmcache.ai/mp/l2_storage.html#mooncake-store-mooncake-store-native-connector), [P2P backend](https://github.com/LMCache/LMCache/pull/4182) | +| FlexKV | TE, Store | Distributed cache reuse and a content-addressed REMOTE tier | [Distributed reuse](https://github.com/taco-project/FlexKV/blob/main/docs/dist_reuse/README_en.md), [Store backend](https://github.com/taco-project/FlexKV/pull/231) | +| NIXL | TE | Mooncake transport plugin for NIXL data movement | [Plugin guide](https://github.com/ai-dynamo/nixl/blob/main/src/plugins/mooncake/README.md) | +| Dynamo | Store via SGLang HiCache | Shared-pool awareness and event-driven KV cache routing | [HiCache integration](https://github.com/ai-dynamo/dynamo/blob/main/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/sglang/hicache.md) | +| llm-d | TE, Store via vLLM | Kubernetes PD serving and tiered KV cache offloading | [PD integration](https://github.com/llm-d/llm-d/pull/1511), [Tiered prefix cache](https://github.com/llm-d/llm-d/tree/main/guides/tiered-prefix-cache) | +| RBG | TE, Store via inference engines | Role-based orchestration of Store, PD transfer, and cache reuse | [Deployment examples](https://github.com/sgl-project/rbg/tree/main/examples/inference/ecosystem/mooncake) | +| srt-slurm | TE, Store via inference engines | Slurm deployment and Mooncake service configuration | [Mooncake services](https://github.com/NVIDIA/srt-slurm/blob/main/docs/mooncake-kv-store.md) | +| NeMo RL | Store via TransferQueue | Rollout-to-training data exchange with optional GDR client transfers | [Data plane](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data_plane/README.md), [GDR support](https://github.com/NVIDIA-NeMo/RL/pull/3501) | +| TransferQueue | Store | RL tensor/object storage, GDR transfers, and SSD offloading | [Store client](https://github.com/Ascend/TransferQueue/blob/main/transfer_queue/storage/clients/mooncake_client.py), [GDR guide](https://github.com/Ascend/TransferQueue/blob/main/docs/storage_backends/mooncake_gdr.md) | +| verl | TE | Checkpoint-engine backend for trainer-to-rollout weight synchronization | [Mooncake backend](https://github.com/verl-project/verl/blob/main/verl/checkpoint_engine/mooncake_checkpoint_engine.py) | +| checkpoint-engine | TE | Peer-to-peer model-weight distribution and updates | [P2P backend](https://github.com/MoonshotAI/checkpoint-engine) | +| Miles | Store | Rollout-to-training data transfer | [Rollout transfer guide](https://github.com/radixark/miles/blob/main/docs/advanced/mooncake-rollout-transfer.md) | +| TorchSpec | Store | Hidden-state exchange between inference and draft-model training | [Training architecture](https://github.com/torchspec-project/TorchSpec) | +| Speculators | Store | Hidden-state exchange for distributed online draft-model training | [Hidden-state connectors](https://github.com/vllm-project/speculators/tree/main/hs_connectors) | +| vLLM-Omni | TE, Store | Data exchange between multimodal pipeline stages | [Omni connectors](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/distributed/omni_connectors/connectors) | +| LightX2V | TE | Encoder/transformer disaggregation for video generation | [Transfer backend](https://github.com/ModelTC/LightX2V/blob/main/lightx2v/disagg/mooncake.py) | +| slime | TE via SGLang | PD-disaggregated rollout serving | [RL serving architecture](https://github.com/THUDM/slime/blob/main/docs/en/advanced/rl-systems.md) | + +The Dynamo entry describes SGLang HiCache shared-cache routing. Its separate [KVBM G4 Store adapter proposal](https://github.com/ai-dynamo/dynamo/pull/9786) is still unmerged. NeMo RL's GDR path changes client-side tensor transfers; queued objects remain in Mooncake-managed host memory. + +
+

πŸ–₯️ Supported Hardware

Mooncake supports hardware backends across accelerator vendors, cloud fabrics, and standard datacenter interconnects, as listed below. See the [supported protocols](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html) and [Transfer Engine design docs](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html) for details. From 19e0ab28f125346f7348e8ac9d043537d7202771 Mon Sep 17 00:00:00 2001 From: lujh <101535776+LujhCoconut@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:09:14 +0800 Subject: [PATCH 03/14] [Store] Fix BatchProbeKey call in master_service DSL scenario (#4729) #4314 added the GrantLeasePolicy parameter to MasterService::BatchProbeKey but missed this call site, which breaks the tent build. Pass the default {} policy: ProbeLeaseMode::None preserves the lease-free existence probe the scenario expects. --- mooncake-store/tests/master_service/dsl/scenario.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mooncake-store/tests/master_service/dsl/scenario.cpp b/mooncake-store/tests/master_service/dsl/scenario.cpp index 850fa72de4..16cd982ff9 100644 --- a/mooncake-store/tests/master_service/dsl/scenario.cpp +++ b/mooncake-store/tests/master_service/dsl/scenario.cpp @@ -1922,8 +1922,8 @@ MasterScenario& MasterScenario::Then(BatchProbeSpec batch_probe) { return *this; } - const auto result = - service_->BatchProbeKey(batch_probe.keys, TenantId(batch_probe.tenant)); + const auto result = service_->BatchProbeKey({}, batch_probe.keys, + TenantId(batch_probe.tenant)); if (result.size() != batch_probe.expected.size()) { Fail("BatchProbe returned " + std::to_string(result.size()) + " results; expected " + From 7ff6bf7f72f15aba7289bc25c87cf491e04efc6f Mon Sep 17 00:00:00 2001 From: tensorcopy Date: Sun, 18 Oct 2026 01:56:03 -0700 Subject: [PATCH 04/14] [TransferEngine] EFA: release the endpoint state lock during handshake RPC (#4704) setupConnectionsByActive() held the endpoint RWSpinlock across the transfer-metadata handshake RPC. The handshake is reciprocal: the peer's passive handler must take the same lock on its own endpoint, so a bilateral active setup (both sides dialing at once) deadlocked. Serialize active setup with a dedicated mutex that passive handlers never take, and hold the state lock only while reading local state and while publishing the reply, so three properties hold: - Active callers are serialized into one handshake RPC; the state lock is released while the RPC is in flight. - A reply that crosses a disconnect, a nic-path change or endpoint teardown is stale: a generation counter, bumped on every disconnect/path change, rejects it instead of publishing a connection against dead state. - A passive handshake that connected the endpoint while the RPC was in flight wins: the active caller returns success and discards its own reply, even if the RPC failed or the generation moved. Also qualify efa_context.h includes as "transport/efa_transport/efa_context.h" so the endpoint header can be included from tests that shadow the context boundary (tests/efa_endpoint_handshake_test.cpp compiles the real endpoint source against a scripted RPC/address-vector fake and runs without an EFA device). --- .../transport/efa_transport/efa_endpoint.h | 8 +- .../transport/efa_transport/efa_endpoint.cpp | 64 +++--- mooncake-transfer-engine/tests/CMakeLists.txt | 21 +- .../transport/efa_transport/efa_context.h | 70 +++++++ .../tests/efa_endpoint_handshake_test.cpp | 194 ++++++++++++++++++ 5 files changed, 332 insertions(+), 25 deletions(-) create mode 100644 mooncake-transfer-engine/tests/efa_endpoint_handshake_fakes/transport/efa_transport/efa_context.h create mode 100644 mooncake-transfer-engine/tests/efa_endpoint_handshake_test.cpp diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h index 550f6b68de..e472a0d3be 100644 --- a/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_endpoint.h @@ -20,11 +20,12 @@ #include #include +#include #include #include #include "common.h" -#include "efa_context.h" +#include "transport/efa_transport/efa_context.h" #include "transfer_metadata.h" #include "transport/transport.h" @@ -117,7 +118,12 @@ class EfaEndPoint { EfaContext& context_; std::atomic status_; + // Only active callers take this mutex. Passive RPC handlers must not. + std::mutex active_setup_mutex_; RWSpinlock lock_; // protects peer_nic_path_ and status_ + // Bumped (under lock_) on every disconnect, nic-path change and teardown, + // so a handshake reply that crosses one of them can be rejected. + uint64_t generation_ = 0; std::string peer_nic_path_; fi_addr_t peer_fi_addr_; // slot in context_.av() // Last peer EFA address successfully inserted into the AV. Used to diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp index 7357a22ffa..6b72196cc6 100644 --- a/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp @@ -36,41 +36,57 @@ void EfaEndPoint::setPeerNicPath(const std::string& peer_nic_path) { << peer_nic_path_ << " -> " << peer_nic_path; disconnectUnlocked(); } + ++generation_; peer_nic_path_ = peer_nic_path; } int EfaEndPoint::setupConnectionsByActive() { - RWSpinlock::WriteGuard guard(lock_); - if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0; - - // Loopback: handshake against ourselves. Use the binary overload - // so we avoid hex-encoding the local address just to have - // insertPeerAddr decode it right back. - if (context_.nicPath() == peer_nic_path_) { - const auto& bytes = context_.localEpAddrBytes(); - int ret = context_.insertPeerAddrBytes(bytes.data(), bytes.size(), - peer_fi_addr_); - if (ret != 0) return ret; - status_.store(CONNECTED, std::memory_order_release); - LOG(INFO) << "EFA loopback connection established: " << toString(); - return 0; - } - - // Exchange addresses via the transfer-metadata handshake RPC. + std::lock_guard setup_guard(active_setup_mutex_); TransferMetadata::HandShakeDesc local_desc, peer_desc; - local_desc.local_nic_path = context_.nicPath(); - local_desc.peer_nic_path = peer_nic_path_; - local_desc.efa_addr = context_.localEpAddr(); + uint64_t generation; + { + RWSpinlock::WriteGuard guard(lock_); + if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0; + + if (context_.nicPath() == peer_nic_path_) { + // Loopback: handshake against ourselves. Use the binary overload + // so we avoid hex-encoding the local address just to have + // insertPeerAddr decode it right back. + const auto& bytes = context_.localEpAddrBytes(); + int ret = context_.insertPeerAddrBytes(bytes.data(), bytes.size(), + peer_fi_addr_); + if (ret != 0) return ret; + status_.store(CONNECTED, std::memory_order_release); + LOG(INFO) << "EFA loopback connection established: " << toString(); + return 0; + } + generation = generation_; + // Exchange addresses via the transfer-metadata handshake RPC. + local_desc.local_nic_path = context_.nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.efa_addr = context_.localEpAddr(); + } - auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); - auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); + auto peer_server_name = getServerNameFromNicPath(local_desc.peer_nic_path); + auto peer_nic_name = getNicNameFromNicPath(local_desc.peer_nic_path); if (peer_server_name.empty() || peer_nic_name.empty()) { - LOG(ERROR) << "Parse peer EFA nic path failed: " << peer_nic_path_; + LOG(ERROR) << "Parse peer EFA nic path failed: " + << local_desc.peer_nic_path; return ERR_INVALID_ARGUMENT; } + // A reciprocal RPC needs our passive handler, which takes lock_. int rc = context_.engine().sendHandshake(peer_server_name, local_desc, peer_desc); + + RWSpinlock::WriteGuard guard(lock_); + // A passive handshake may have connected the endpoint while the RPC was + // in flight; reuse that connection and discard this reply. + if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0; + if (generation != generation_) { + LOG(WARNING) << "Discarding stale EFA handshake reply: " << toString(); + return ERR_ENDPOINT; + } if (rc) return rc; if (peer_desc.efa_addr.empty()) { @@ -157,6 +173,7 @@ void EfaEndPoint::disconnect() { } void EfaEndPoint::disconnectUnlocked() { + ++generation_; if (peer_fi_addr_ != FI_ADDR_UNSPEC) { context_.removePeerAddr(peer_fi_addr_); peer_fi_addr_ = FI_ADDR_UNSPEC; @@ -167,6 +184,7 @@ void EfaEndPoint::disconnectUnlocked() { void EfaEndPoint::markDetachedForTeardown() { RWSpinlock::WriteGuard guard(lock_); + ++generation_; peer_fi_addr_ = FI_ADDR_UNSPEC; cached_peer_addr_.clear(); status_.store(UNCONNECTED, std::memory_order_release); diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index aa29927ec3..3a865cdfea 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -313,9 +313,28 @@ if(USE_EFA) add_executable(efa_select_device_test ${WORKSPACE}/efa_select_device_test.cpp) target_link_libraries(efa_select_device_test PUBLIC transfer_engine gtest - gtest_main) + gtest_main) add_test(NAME efa_select_device_test COMMAND efa_select_device_test) + # Compiles the real EfaEndPoint source against the test-only + # RPC/address-vector fakes in efa_endpoint_handshake_fakes/ (shadowed + # efa_context.h via -iquote), so the endpoint locking and generation logic + # runs without an EFA device. Bounded: a locking regression must fail, not + # hang the test runner. + add_executable( + efa_endpoint_handshake_test + ${WORKSPACE}/efa_endpoint_handshake_test.cpp + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/src/transport/efa_transport/efa_endpoint.cpp + ) + target_compile_options( + efa_endpoint_handshake_test + PRIVATE -iquote${WORKSPACE}/efa_endpoint_handshake_fakes) + target_link_libraries( + efa_endpoint_handshake_test PUBLIC glog::glog gflags::gflags gtest + gtest_main pthread) + add_test(NAME efa_endpoint_handshake_test COMMAND efa_endpoint_handshake_test) + set_tests_properties(efa_endpoint_handshake_test PROPERTIES TIMEOUT 120) + add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp) target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME efa_c_api_test COMMAND efa_c_api_test) diff --git a/mooncake-transfer-engine/tests/efa_endpoint_handshake_fakes/transport/efa_transport/efa_context.h b/mooncake-transfer-engine/tests/efa_endpoint_handshake_fakes/transport/efa_transport/efa_context.h new file mode 100644 index 0000000000..9a75cdace8 --- /dev/null +++ b/mooncake-transfer-engine/tests/efa_endpoint_handshake_fakes/transport/efa_transport/efa_context.h @@ -0,0 +1,70 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Test-only RPC and address-vector boundary for the real EfaEndPoint code. +// This header shadows "transport/efa_transport/efa_context.h" on the test's +// include path, so efa_endpoint.cpp can be compiled and driven without an +// EFA device or a real provider: the handshake RPC and the AV insert/remove +// calls are recorded and scripted by the test. +#pragma once + +#include +#include +#include + +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +struct EfaHandshakeTestTransport { + using Desc = TransferMetadata::HandShakeDesc; + std::function rpc; + int sendHandshake(const std::string&, const Desc& local, Desc& peer) { + return rpc(local, peer); + } +}; + +class EfaContext { + public: + std::string name; + EfaHandshakeTestTransport transport; + int inserts = 0; + int removals = 0; + explicit EfaContext(std::string path) : name(std::move(path)) {} + const std::string& nicPath() const { return name; } + std::string localEpAddr() const { return name; } + const std::vector& localEpAddrBytes() const { + static const std::vector bytes{1}; + return bytes; + } + EfaHandshakeTestTransport& engine() { return transport; } + int insertPeerAddr(const std::string&, fi_addr_t& addr) { + addr = ++inserts; + return 0; + } + int insertPeerAddrBytes(const void*, size_t, fi_addr_t& addr) { + return insertPeerAddr("", addr); + } + int removePeerAddr(fi_addr_t) { + ++removals; + return 0; + } + int submitSlicesOnPeer(fi_addr_t, std::vector&, + std::vector&) { + return 0; + } +}; + +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/efa_endpoint_handshake_test.cpp b/mooncake-transfer-engine/tests/efa_endpoint_handshake_test.cpp new file mode 100644 index 0000000000..4776595e3c --- /dev/null +++ b/mooncake-transfer-engine/tests/efa_endpoint_handshake_test.cpp @@ -0,0 +1,194 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Behavioral tests for the EFA endpoint handshake state machine. The real +// EfaEndPoint source is compiled against the test-only RPC/address-vector +// fakes in efa_endpoint_handshake_fakes/ (a shadowed efa_context.h), so +// the locking and generation logic runs without an EFA device. All waits +// are bounded: a locking regression must fail, not hang the runner. +#include +#include +#include + +#include + +#include "transport/efa_transport/efa_endpoint.h" + +namespace mooncake { + +namespace { + +using Desc = TransferMetadata::HandShakeDesc; + +Desc Reply(const Desc& request) { + Desc result; + result.local_nic_path = request.peer_nic_path; + result.peer_nic_path = request.local_nic_path; + result.efa_addr = result.local_nic_path; + return result; +} + +TEST(EfaEndPointHandshakeTest, ReciprocalActiveSetupDoesNotDeadlock) { + constexpr auto kWait = std::chrono::seconds(30); + EfaContext a("a:1@efa"), b("b:2@efa"); + EfaEndPoint ea(a), eb(b); + ea.setPeerNicPath(b.name); + eb.setPeerNicPath(a.name); + std::promise a_started, b_started; + auto a_ready = a_started.get_future(); + auto b_ready = b_started.get_future(); + a.transport.rpc = [&](const Desc& request, Desc& response) { + a_started.set_value(); + if (b_ready.wait_for(kWait) != std::future_status::ready) { + return ERR_SOCKET; + } + return eb.setupConnectionsByPassive(request, response); + }; + b.transport.rpc = [&](const Desc& request, Desc& response) { + b_started.set_value(); + if (a_ready.wait_for(kWait) != std::future_status::ready) { + return ERR_SOCKET; + } + return ea.setupConnectionsByPassive(request, response); + }; + // Each side's passive handler runs while the other side's active setup + // is blocked in its handshake RPC. This only completes if the active + // path does not hold the endpoint state lock across the RPC. + auto first = std::async(std::launch::async, + [&] { return ea.setupConnectionsByActive(); }); + auto second = std::async(std::launch::async, + [&] { return eb.setupConnectionsByActive(); }); + ASSERT_EQ(std::future_status::ready, first.wait_for(kWait)) + << "reciprocal active setup deadlocked"; + ASSERT_EQ(std::future_status::ready, second.wait_for(kWait)) + << "reciprocal active setup deadlocked"; + ASSERT_EQ(first.get(), 0); + ASSERT_EQ(second.get(), 0); + EXPECT_TRUE(ea.connected() && eb.connected()); + // Depending on which side connects first, a late passive handshake may + // re-insert the peer address. Either way each side ends with exactly one + // live address. + EXPECT_EQ(1, a.inserts - a.removals); + EXPECT_EQ(1, b.inserts - b.removals); + EXPECT_LE(a.removals, 1); + EXPECT_LE(b.removals, 1); +} + +void CancelledHandshake(bool change_path, bool detach) { + constexpr auto kWait = std::chrono::seconds(30); + EfaContext context("a:1@efa"); + EfaEndPoint endpoint(context); + endpoint.setPeerNicPath("b:2@efa"); + std::promise started, release; + auto ready = started.get_future(); + auto released = release.get_future(); + context.transport.rpc = [&](const Desc& request, Desc& response) { + started.set_value(); + if (released.wait_for(kWait) != std::future_status::ready) { + return ERR_SOCKET; + } + response = Reply(request); + return 0; + }; + auto active = std::async(std::launch::async, [&] { + return endpoint.setupConnectionsByActive(); + }); + ASSERT_EQ(std::future_status::ready, ready.wait_for(kWait)); + if (change_path) { + endpoint.setPeerNicPath("c:3@efa"); + } else if (detach) { + endpoint.markDetachedForTeardown(); + } else { + endpoint.disconnect(); + } + release.set_value(); + // The reply crosses a disconnect/nic-path change/teardown and is stale. + ASSERT_EQ(std::future_status::ready, active.wait_for(kWait)); + EXPECT_NE(active.get(), 0); + EXPECT_FALSE(endpoint.connected()); + EXPECT_EQ(0, context.inserts); +} + +TEST(EfaEndPointHandshakeTest, StaleReplyAfterDisconnectIsRejected) { + CancelledHandshake(/*change_path=*/false, /*detach=*/false); +} + +TEST(EfaEndPointHandshakeTest, StaleReplyAfterNicPathChangeIsRejected) { + CancelledHandshake(/*change_path=*/true, /*detach=*/false); +} + +TEST(EfaEndPointHandshakeTest, StaleReplyAfterTeardownIsRejected) { + CancelledHandshake(/*change_path=*/false, /*detach=*/true); +} + +TEST(EfaEndPointHandshakeTest, ConcurrentActiveCallersShareOneHandshake) { + EfaContext context("a:1@efa"); + EfaEndPoint endpoint(context); + endpoint.setPeerNicPath("b:2@efa"); + int calls = 0; + context.transport.rpc = [&](const Desc& request, Desc& response) { + ++calls; + response = Reply(request); + return 0; + }; + std::vector> callers; + for (int i = 0; i != 8; ++i) { + callers.push_back(std::async(std::launch::async, [&] { + return endpoint.setupConnectionsByActive(); + })); + } + for (auto& call : callers) { + ASSERT_EQ(call.get(), 0); + } + EXPECT_EQ(1, calls); + EXPECT_EQ(1, context.inserts); + endpoint.disconnect(); + context.transport.rpc = [](const Desc&, Desc&) { return ERR_SOCKET; }; + EXPECT_EQ(ERR_SOCKET, endpoint.setupConnectionsByActive()); + EXPECT_FALSE(endpoint.connected()); + context.transport.rpc = [&](const Desc& request, Desc& response) { + response = Reply(request); + return 0; + }; + EXPECT_EQ(0, endpoint.setupConnectionsByActive()); +} + +TEST(EfaEndPointHandshakeTest, + ActiveSetupReusesPassiveConnectionAfterRpcFailure) { + EfaContext context("a:1@efa"); + EfaEndPoint endpoint(context); + endpoint.setPeerNicPath("b:2@efa"); + context.transport.rpc = [&](const Desc& request, Desc&) { + Desc incoming = Reply(request), response; + EXPECT_EQ(0, endpoint.setupConnectionsByPassive(incoming, response)); + // A repeated passive handshake re-inserts the peer address (see the + // comment in setupConnectionsByPassive). The in-flight active setup + // must still succeed on the resulting connection. + EXPECT_EQ(0, endpoint.setupConnectionsByPassive(incoming, response)); + return ERR_SOCKET; + }; + EXPECT_EQ(0, endpoint.setupConnectionsByActive()); + EXPECT_TRUE(endpoint.connected()); + EXPECT_EQ(2, context.inserts); + EXPECT_EQ(1, context.removals); + endpoint.disconnect(); + context.transport.rpc = [](const Desc&, Desc&) { return 0; }; + // A reciprocal reply with an empty/invalid address is rejected. + EXPECT_EQ(ERR_REJECT_HANDSHAKE, endpoint.setupConnectionsByActive()); + EXPECT_FALSE(endpoint.connected()); + EXPECT_EQ(2, context.inserts); +} + +} // namespace +} // namespace mooncake From 5d0636f5d39f9de729a58ea036f98c395dd0b291 Mon Sep 17 00:00:00 2001 From: ykwd Date: Mon, 19 Oct 2026 00:34:33 +0800 Subject: [PATCH 05/14] [Doc] Add a blogs section and hide duplicated nav on the docs landing page (#4693) --------- Co-authored-by: Ke Yang Co-authored-by: Cursor --- docs/source/_static/custom.css | 28 +++++++++++++++ docs/source/index.md | 64 ++++++++++++++++++++++++++++------ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/docs/source/_static/custom.css b/docs/source/_static/custom.css index 79bd2082b4..3d63e2798f 100644 --- a/docs/source/_static/custom.css +++ b/docs/source/_static/custom.css @@ -6,3 +6,31 @@ margin: 0; } } + +.latest-blogs .sd-card { + height: 100%; + border-radius: 0.75rem; + transition: border-color 0.15s ease, box-shadow 0.15s ease; +} + +.latest-blogs .sd-card:hover { + border-color: var(--pst-color-primary); + box-shadow: 0 4px 14px color-mix(in srgb, var(--pst-color-text-base) 8%, transparent); +} + +.latest-blogs .sd-card-img-top { + width: 100%; + height: 148px; + object-fit: cover; +} + +.latest-blogs .sd-card-title { + font-size: 1rem; + line-height: 1.4; +} + +.latest-blogs .blog-card-date p { + margin-bottom: 0; + color: var(--pst-color-text-muted); + font-size: 0.875rem; +} diff --git a/docs/source/index.md b/docs/source/index.md index 23fbf5be46..61c61ddd36 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -82,13 +82,60 @@ At the center of Mooncake is a KVCache-centric scheduler that balances effective ::: -## Documentation +

πŸ“° Blogs

+ +::::{grid} 1 2 2 2 +:gutter: 2 +:class-container: latest-blogs + +:::{grid-item-card} Over One Trillion Tokens per Day: How Mooncake Powers Approaching.AI’s Token Factory +:img-top: https://kvcache.ai/blog/mooncake-in-approaching-ai/featured_hue86b549ecdd836b0c6b3aa760df33338_1244829_2a2b63e1a97ff4921bf0e665a65e1403.webp +:img-alt: The HiCache + Mooncake Store Architecture within a PD Group +:link: https://kvcache.ai/blog/mooncake-in-approaching-ai/ +:link-alt: Over One Trillion Tokens per Day: How Mooncake Powers Approaching.AI’s Token Factory +:class-body: blog-card-date + +September 15, 2026 +::: + +:::{grid-item-card} Mooncake for Miles: From Fragmented Rollout Data to Efficient Bulk I/O +:img-top: https://kvcache.ai/blog/mooncake-rl-rollout-data-transfer/featured_hue655d0112fa862027ab204efc8ff95c7_58042_8bb94c00d58ac807a3ff4bc8f8df21e2.webp +:img-alt: Miles rollout data transfer with Mooncake +:link: https://kvcache.ai/blog/mooncake-rl-rollout-data-transfer/ +:link-alt: Mooncake for Miles: From Fragmented Rollout Data to Efficient Bulk I/O +:class-body: blog-card-date + +August 20, 2026 +::: + +:::{grid-item-card} Scaling KV Cache Beyond Memory with Mooncake SSD Offloading +:img-top: https://kvcache.ai/blog/scaling-kv-cache-beyond-memory/featured_huc52d82bb227581a966940933903d5e2c_479035_349926621d7fc55e030b6d364f46602b.webp +:img-alt: Mooncake SSD Offloading Architecture +:link: https://kvcache.ai/blog/scaling-kv-cache-beyond-memory/ +:link-alt: Scaling KV Cache Beyond Memory with Mooncake SSD Offloading +:class-body: blog-card-date + +July 15, 2026 +::: + +:::{grid-item-card} How Much KV Cache Budget Do We Need for LLM Serving? +:img-top: https://kvcache.ai/blog/calculate-kvcache-cache-budge/featured_hu137d5fd64784a367d2cae0088c081616_54073_4242cb61033db09c45fd8739c71f32e2.webp +:img-alt: KV Cache Hit Rate Figure +:link: https://kvcache.ai/blog/calculate-kvcache-cache-budge/ +:link-alt: How Much KV Cache Budget Do We Need for LLM Serving? +:class-body: blog-card-date + +June 26, 2026 +::: + +:::: % How to start using Mooncake? :::{toctree} :caption: Getting Started :maxdepth: 1 +:hidden: getting_started/quick-start getting_started/build @@ -102,6 +149,7 @@ getting_started/store-rs-features :::{toctree} :caption: Deployment :maxdepth: 1 +:hidden: deployment/mooncake-store-deployment-guide deployment/store-rs/index @@ -118,6 +166,7 @@ Mooncake x LMDeploy Integration :::{toctree} :caption: Performance :maxdepth: 1 +:hidden: performance/mooncake/index performance/sglang/index performance/vllm/index @@ -128,6 +177,7 @@ performance/vllm/index :::{toctree} :caption: Design Documents :maxdepth: 1 +:hidden: design/architecture design/transfer-engine/index @@ -147,6 +197,7 @@ design/hicache-design :::{toctree} :caption: API Reference :maxdepth: 1 +:hidden: api-reference/python/index api-reference/cpp/index @@ -159,6 +210,7 @@ api-reference/rust/index :::{toctree} :caption: Troubleshooting :maxdepth: 1 +:hidden: troubleshooting/error-code troubleshooting/troubleshooting @@ -172,15 +224,7 @@ troubleshooting/store-rs-key-access :::{toctree} :caption: Community :maxdepth: 1 +:hidden: community/governance ::: - -% Archived content - -### Archived - -- [Guide: vLLM MooncakeStoreConnector](deployment/integrations/vllm/vllm-mooncakestoreconnector) -- [vLLM V0 Disaggregated Serving Demo](deployment/integrations/vllm/vllm-integration-v0.2) -- [vLLM V0 Disaggregated Serving with MooncakeStore](deployment/integrations/vllm/vllm-integration-v0.3) -- [vLLM v1 backend Disaggregated Serving with MooncakeConnector](deployment/integrations/vllm/vllm-integration-v1.0) From 2264425fdb2a78c9a424d15f1ed532e9823b67f4 Mon Sep 17 00:00:00 2001 From: ykwd Date: Mon, 19 Oct 2026 10:10:30 +0800 Subject: [PATCH 06/14] [Doc] Fold the Store-RS build guide and feature page into the docs layout (#4689) * Move feature page to store-rs * Move store-rs build guide to the build file * Highlight the title of store-rs toubleshooting * Address review comments --------- Co-authored-by: Ke Yang --- .../api-reference/python/mooncake-store.md | 2 +- docs/source/api-reference/python/store-rs.md | 2 +- docs/source/deployment/store-rs/index.md | 2 +- docs/source/deployment/store-rs/python.md | 4 +- .../deployment/store-rs/rolling-upgrade.md | 2 +- .../deployment/store-rs/route-migration.md | 2 +- .../design/store/store-rs/components.md | 4 +- .../store/store-rs/features.md} | 4 +- docs/source/design/store/store-rs/index.md | 2 + docs/source/getting_started/build.md | 64 ++++++++++++++ docs/source/getting_started/store-rs.md | 88 ------------------- docs/source/index.md | 2 - .../mooncake/store-rs-benchmark.md | 2 +- .../troubleshooting/store-rs-key-access.md | 2 +- mooncake-store-rs/README.md | 4 +- 15 files changed, 81 insertions(+), 105 deletions(-) rename docs/source/{getting_started/store-rs-features.md => design/store/store-rs/features.md} (99%) delete mode 100644 docs/source/getting_started/store-rs.md diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md index b860821ca8..3c9eb87c43 100644 --- a/docs/source/api-reference/python/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -30,7 +30,7 @@ imported and accepts cpp or rs. This page describes the shared facade and the default C++ backend. See the [Store-RS Python API](store-rs.md) for backend-specific runtime behavior, API coverage, and compatibility notes. The -[source-checkout quickstart](../../getting_started/store-rs.md) shows how to +{ref}`Build with Store-RS ` section shows how to build the root wheel with Store-RS enabled. ## Quick Start diff --git a/docs/source/api-reference/python/store-rs.md b/docs/source/api-reference/python/store-rs.md index b16cc170a5..56aeb1009c 100644 --- a/docs/source/api-reference/python/store-rs.md +++ b/docs/source/api-reference/python/store-rs.md @@ -4,7 +4,7 @@ The Python Store facade selects its implementation when mooncake.store is first imported. The shared API and default C++ backend are documented in [Mooncake Store Python API](mooncake-store.md). This page describes Store-RS runtime and API behavior. See the -[source-checkout quickstart](../../getting_started/store-rs.md) to build and +{ref}`Build with Store-RS ` section to build and install the unified root wheel. ## What the Python Layer Provides diff --git a/docs/source/deployment/store-rs/index.md b/docs/source/deployment/store-rs/index.md index 90dfec28ff..c35a4f4541 100644 --- a/docs/source/deployment/store-rs/index.md +++ b/docs/source/deployment/store-rs/index.md @@ -2,7 +2,7 @@ Use these guides to build the optional Store-RS component, configure a runtime, operate tenant policy and routes, and validate deployment scenarios. The -[source-checkout quickstart](../../getting_started/store-rs.md) starts with a +{ref}`Build with Store-RS ` section starts with a root CMake build and the unified Python wheel. | Guide | Use it to | diff --git a/docs/source/deployment/store-rs/python.md b/docs/source/deployment/store-rs/python.md index 334a48009e..242eb45789 100644 --- a/docs/source/deployment/store-rs/python.md +++ b/docs/source/deployment/store-rs/python.md @@ -1,11 +1,11 @@ # Store-RS Python Operations The root wheel installs the standalone Store-RS commands. Build and install it -using the [source-checkout quickstart](../../getting_started/store-rs.md). +using {ref}`Build with Store-RS `. ## Standalone Store-RS Commands -The root CMake target build_store_rs_cli builds the client, admin, and benchmark binaries. The unified root wheel installs the standalone commands under their mooncake-store-rs-* names. Use the [source-checkout quickstart](../../getting_started/store-rs.md) to configure the root build and install the wheel. +The root CMake target build_store_rs_cli builds the client, admin, and benchmark binaries. The unified root wheel installs the standalone commands under their mooncake-store-rs-* names. Use {ref}`Build with Store-RS ` to configure the root build and install the wheel. Start a storage client: diff --git a/docs/source/deployment/store-rs/rolling-upgrade.md b/docs/source/deployment/store-rs/rolling-upgrade.md index e6984a8023..946cc43b15 100644 --- a/docs/source/deployment/store-rs/rolling-upgrade.md +++ b/docs/source/deployment/store-rs/rolling-upgrade.md @@ -8,7 +8,7 @@ This document describes how to perform a rolling upgrade on a Mooncake Store clu - The new version only increments `store_api_minor_version` (minor version) and does **not** change `store_api_version` (major version) - The changes in the new version are backward-compatible (e.g., adding fields, adding RPCs) and do not alter existing API semantics -- Both versions were built from their own root CMake build trees using the [Store-RS source-checkout quickstart](../../getting_started/store-rs.md) +- Both versions were built from their own root CMake build trees using {ref}`Build with Store-RS ` ### Compatibility Guarantees diff --git a/docs/source/deployment/store-rs/route-migration.md b/docs/source/deployment/store-rs/route-migration.md index 74c5dd5835..b15902a505 100644 --- a/docs/source/deployment/store-rs/route-migration.md +++ b/docs/source/deployment/store-rs/route-migration.md @@ -202,6 +202,6 @@ What this feature is not: ## 9. Related Docs - [Python Guide](../../api-reference/python/store-rs.md) -- [Features](../../getting_started/store-rs-features.md) +- [Features](../../design/store/store-rs/features.md) - [Configuration Reference](configuration.md) - [Architecture](../../design/store/store-rs/architecture.md) diff --git a/docs/source/design/store/store-rs/components.md b/docs/source/design/store/store-rs/components.md index 022f304781..a4d441288b 100644 --- a/docs/source/design/store/store-rs/components.md +++ b/docs/source/design/store/store-rs/components.md @@ -231,11 +231,11 @@ See the [benchmark guide](../../../performance/mooncake/store-rs-benchmark.md) f ## Suggested Reading Order -1. [Source-checkout quickstart](../../../getting_started/store-rs.md) +1. {ref}`Build with Store-RS ` 2. [Deployment](../../../deployment/store-rs/deployment.md) 3. [Rust API](../../../api-reference/rust/store-rs.md) 4. [Configuration](../../../deployment/store-rs/configuration.md) 5. `crates/mooncake-store-core` -6. [Features](../../../getting_started/store-rs-features.md) +6. [Features](features.md) 7. [Architecture](architecture.md) 8. [Test and validation guide](../../../deployment/store-rs/testing.md) diff --git a/docs/source/getting_started/store-rs-features.md b/docs/source/design/store/store-rs/features.md similarity index 99% rename from docs/source/getting_started/store-rs-features.md rename to docs/source/design/store/store-rs/features.md index c085ad3a0a..ac49836db7 100644 --- a/docs/source/getting_started/store-rs-features.md +++ b/docs/source/design/store/store-rs/features.md @@ -149,7 +149,7 @@ For simpler bring-up or debugging, the runtime can store and resolve routes dire The current implementation also exposes admin-driven explicit route migration tasks for key-level `copy` and `move`. -See [Route Migration Guide](../deployment/store-rs/route-migration.md) for the current +See [Route Migration Guide](../../../deployment/store-rs/route-migration.md) for the current operator workflow, task semantics, and request examples. What it provides: @@ -507,7 +507,7 @@ Both paths are covered by repository validation scripts. The root validation entry point is `scripts/ci/run_store_rs_smoke.sh`. It covers the installed Python facade and the dummy and routed TCP/classic-TE -read/write paths. See the [validation guide](../deployment/store-rs/testing.md) +read/write paths. See the [validation guide](../../../deployment/store-rs/testing.md) for required build inputs and direct commands for additional scenarios. Manual path-specific validators stay available as: diff --git a/docs/source/design/store/store-rs/index.md b/docs/source/design/store/store-rs/index.md index 016f864d6c..2d8e5e082e 100644 --- a/docs/source/design/store/store-rs/index.md +++ b/docs/source/design/store/store-rs/index.md @@ -13,6 +13,7 @@ public language APIs live under the Rust and Python API references. | NoF integration boundaries | [NoF](nof) | | Tenant control plane | [Multi-tenant admin design](multi-tenant-admin-control-plane) | | Strict quota accounting | [Tenant quota consistency](tenant-quota-consistency) | +| Implemented capabilities | [Features](features) | :::{toctree} :maxdepth: 1 @@ -23,4 +24,5 @@ cold-tier-design nof multi-tenant-admin-control-plane tenant-quota-consistency +features ::: diff --git a/docs/source/getting_started/build.md b/docs/source/getting_started/build.md index ec6a4deb10..371de0f038 100644 --- a/docs/source/getting_started/build.md +++ b/docs/source/getting_started/build.md @@ -75,6 +75,69 @@ sudo make install `-DUSE_NOF=ON` builds the NoF registration APIs and deployment tools. Use `-DUSE_NOF=OFF` or omit the option when the NVMe-oF SSD pool is not needed. +(build-with-store-rs)= +### Build with Store-RS + +Store-RS is an optional Rust-native Mooncake Store. The default CMake build +leaves it disabled, and the Python Store facade uses the C++ backend. Enable +it when building a root wheel that should include the Rust backend and its standalone commands. + +The build needs a root checkout with submodules, CMake, a C++ toolchain, Rust +(`cargo` on `PATH`), and Python 3.10 or newer. A multi-client runtime also +needs Redis or etcd. + +```bash +export MOONCAKE_ROOT_DIR="$PWD" +export MOONCAKE_STORE_RS_DIR="$MOONCAKE_ROOT_DIR/mooncake-store-rs" +export MOONCAKE_BUILD_DIR="$MOONCAKE_ROOT_DIR/build/store-rs" +export MOONCAKE_PYTHON_BIN="$(command -v python3)" + +cmake -S "$MOONCAKE_ROOT_DIR" -B "$MOONCAKE_BUILD_DIR" \ + -DPython3_EXECUTABLE="$MOONCAKE_PYTHON_BIN" \ + -DWITH_STORE=ON \ + -DWITH_STORE_RUST=OFF \ + -DWITH_STORE_RS=ON \ + -DWITH_TE=ON \ + -DUSE_TENT=ON \ + -DBUILD_SHARED_LIBS=ON \ + -DCMAKE_BUILD_TYPE=Release +cmake --build "$MOONCAKE_BUILD_DIR" --target build_store_rs --parallel +``` + +This out-of-tree build provides the native artifacts and standalone binaries. For a Python wheel only, skip the separate CMake build command; the wheel builder configures and builds its own CMake tree. + +The root wheel owns the mooncake Python package. Build and install one wheel, +then select its Store backend before importing mooncake.store: + +| Setting | Backend | +|---------|---------| +| Unset or cpp | C++ Store | +| rs | Store-RS, included by WITH_STORE_RS=ON | + +The facade reads MOONCAKE_STORE_BACKEND when mooncake.store is first imported. +Use exactly cpp or rs. The setting does not change the standalone command +names. + +```bash +MOONCAKE_STORE_BACKEND=rs .venv-store-rs/bin/python -c \ + "from mooncake.store import MooncakeDistributedStore; print(MooncakeDistributedStore)" + +.venv-store-rs/bin/mooncake-store-rs-client --help +.venv-store-rs/bin/mooncake-store-rs-admin --help +.venv-store-rs/bin/mooncake-store-rs-bench --help +``` + +For workspace unit tests and installed-package smoke coverage, run +`scripts/ci/run_store_rs_smoke.sh` after installing the CMake `python` +component. The [Store-RS validation guide](../deployment/store-rs/testing.md#local-validation) +lists the required CMake paths and the direct commands for heavier scenarios. + +Configuration, validation, and API details are in the +[Store-RS deployment guides](../deployment/store-rs/index.md), the +[Rust API](../api-reference/rust/store-rs.md), the +[Python API](../api-reference/python/store-rs.md), and the +[Admin HTTP API](../api-reference/http/store-rs-admin.md). + ### RISC-V Build Mooncake supports native 64-bit RISC-V Linux builds with `USE_RISCV` enabled. @@ -301,6 +364,7 @@ The following options can be passed to `cmake ..`. | `-DWITH_P2P_STORE=ON/OFF` | `OFF` | Enable Golang support and build the P2P Store component. Requires Go 1.23+. | | `-DWITH_RUST_EXAMPLE=ON/OFF` | `OFF` | Build the Transfer Engine Rust library (`transfer_engine_rust`), tests, and sample benchmark. | | `-DWITH_STORE_RUST=ON/OFF` | `ON` | Build Mooncake Store Rust bindings and CMake Rust targets. | +| `-DWITH_STORE_RS=ON/OFF` | `OFF` | Build the Rust-native Mooncake Store (Store-RS), its Python extension, and the standalone commands. Requires `-DWITH_TE=ON`, `-DUSE_TENT=ON`, and `cargo`. See {ref}`Build with Store-RS `. | | `-DWITH_EP=ON/OFF` | `OFF` | Build Mooncake EP and the native Mooncake PG core/device libraries. The Torch-facing PG adapter is compiled by the installed package at first `import mooncake.pg`. | ### Build Behavior Options diff --git a/docs/source/getting_started/store-rs.md b/docs/source/getting_started/store-rs.md deleted file mode 100644 index 962042e703..0000000000 --- a/docs/source/getting_started/store-rs.md +++ /dev/null @@ -1,88 +0,0 @@ -# Store-RS Source-Checkout Quickstart - -Store-RS is an optional Rust-native Mooncake Store component in the root -repository. The normal CMake build leaves Store-RS disabled, and the Python -Store facade uses the C++ backend by default. Enable Store-RS when building a -root wheel that should include the Rust backend and its standalone commands. - -## Requirements - -- A Mooncake root checkout with its submodules initialized. -- CMake, a C++ toolchain, and the Rust toolchain used by the Store-RS build. -- Python 3.10 or newer for the Python wheel. -- A Redis or etcd metadata service for a multi-client runtime. - -## Configure and build - -Run these commands from the Mooncake repository root. Keep the source, build, -and Python interpreter paths explicit so CMake, the wheel, and validation -scripts use the same checkout: - -```bash -export MOONCAKE_ROOT_DIR="$PWD" -export MOONCAKE_STORE_RS_DIR="$MOONCAKE_ROOT_DIR/mooncake-store-rs" -export MOONCAKE_BUILD_DIR="$MOONCAKE_ROOT_DIR/build/store-rs" -export MOONCAKE_PYTHON_BIN="$(command -v python3)" - -cmake -S "$MOONCAKE_ROOT_DIR" -B "$MOONCAKE_BUILD_DIR" \ - -DPython3_EXECUTABLE="$MOONCAKE_PYTHON_BIN" \ - -DWITH_STORE=ON \ - -DWITH_STORE_RUST=OFF \ - -DWITH_STORE_RS=ON \ - -DWITH_TE=ON \ - -DUSE_TENT=ON \ - -DBUILD_SHARED_LIBS=ON \ - -DCMAKE_BUILD_TYPE=Release -cmake --build "$MOONCAKE_BUILD_DIR" --target build_store_rs --parallel -``` - -This out-of-tree build provides the native artifacts and standalone binaries. -For a Python wheel only, skip the separate CMake build command; the wheel -builder configures and builds its own CMake tree. - -The root Python wheel packages the C++ Store facade and, when configured with -WITH_STORE_RS=ON, the Store-RS extension and the three standalone commands: - -```bash -SKBUILD_CMAKE_ARGS="-DWITH_STORE=ON -DWITH_STORE_RUST=OFF \ - -DWITH_STORE_RS=ON -DWITH_TE=ON -DUSE_TENT=ON \ - -DBUILD_SHARED_LIBS=ON -DPython3_EXECUTABLE=$MOONCAKE_PYTHON_BIN" \ - "$MOONCAKE_PYTHON_BIN" -m pip wheel . --no-deps --wheel-dir dist/store-rs - -"$MOONCAKE_PYTHON_BIN" -m venv .venv-store-rs -.venv-store-rs/bin/python -m pip install dist/store-rs/mooncake_transfer_engine-*.whl -``` - -The root wheel owns the mooncake Python package. Build and install one wheel, -then select its Store backend before importing mooncake.store: - -| Setting | Backend | -|---------|---------| -| Unset or cpp | C++ Store | -| rs | Store-RS, included by WITH_STORE_RS=ON | - -The facade reads MOONCAKE_STORE_BACKEND when mooncake.store is first imported. -Use exactly cpp or rs. The setting does not change the standalone command -names. - -```bash -MOONCAKE_STORE_BACKEND=rs .venv-store-rs/bin/python -c \ - "from mooncake.store import MooncakeDistributedStore; print(MooncakeDistributedStore)" - -.venv-store-rs/bin/mooncake-store-rs-client --help -.venv-store-rs/bin/mooncake-store-rs-admin --help -.venv-store-rs/bin/mooncake-store-rs-bench --help -``` - -## Validate and continue - -For workspace unit tests and installed-package smoke coverage, run -`scripts/ci/run_store_rs_smoke.sh` after installing the CMake `python` -component. The [Store-RS validation guide](../deployment/store-rs/testing.md#local-validation) -lists the required CMake paths and the direct commands for heavier scenarios. - -The [Store-RS deployment index](../deployment/store-rs/index.md) links to -configuration, operator, and scenario-specific validation guides. For API -details, see the [Rust API](../api-reference/rust/store-rs.md), the -[Python API](../api-reference/python/store-rs.md), and the -[Admin HTTP API](../api-reference/http/store-rs-admin.md). diff --git a/docs/source/index.md b/docs/source/index.md index 61c61ddd36..f65560c708 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -139,8 +139,6 @@ June 26, 2026 getting_started/quick-start getting_started/build -getting_started/store-rs -getting_started/store-rs-features ::: diff --git a/docs/source/performance/mooncake/store-rs-benchmark.md b/docs/source/performance/mooncake/store-rs-benchmark.md index 8ae6ff77b9..262cd2150d 100644 --- a/docs/source/performance/mooncake/store-rs-benchmark.md +++ b/docs/source/performance/mooncake/store-rs-benchmark.md @@ -275,7 +275,7 @@ Keyspace behavior is conditional: ## Build -Configure the root CMake project as described in the [Store-RS source-checkout quickstart](../../getting_started/store-rs.md), then build the build_store_rs_cli target. The unified root wheel installs the client, admin, and benchmark commands for direct use from its Python environment. +Configure the root CMake project as described in {ref}`Build with Store-RS `, then build the build_store_rs_cli target. The unified root wheel installs the client, admin, and benchmark commands for direct use from its Python environment. ## Next reading diff --git a/docs/source/troubleshooting/store-rs-key-access.md b/docs/source/troubleshooting/store-rs-key-access.md index 049e1b243f..9fbe391b7b 100644 --- a/docs/source/troubleshooting/store-rs-key-access.md +++ b/docs/source/troubleshooting/store-rs-key-access.md @@ -1,4 +1,4 @@ -# Key Access Path Debugging +# Store-RS Troubleshooting This guide explains how to trace a single key through Mooncake Store RS: where a `put` writes the object, how a `get` resolves and reads it, which route layer is diff --git a/mooncake-store-rs/README.md b/mooncake-store-rs/README.md index 15b021e56e..6b0ad4909c 100644 --- a/mooncake-store-rs/README.md +++ b/mooncake-store-rs/README.md @@ -95,8 +95,8 @@ Python facade and does not change the standalone command behavior. | Topic | Guide | |-------|-------| -| Source-checkout build and wheel | [Store-RS quickstart](../docs/source/getting_started/store-rs.md) | -| Feature overview | [Store-RS features](../docs/source/getting_started/store-rs-features.md) | +| Source-checkout build and wheel | [Build with Store-RS](../docs/source/getting_started/build.md#build-with-store-rs) | +| Feature overview | [Store-RS features](../docs/source/design/store/store-rs/features.md) | | Rust integration | [Rust API](../docs/source/api-reference/rust/store-rs.md) | | Python facade and Store-RS semantics | [Python API](../docs/source/api-reference/python/store-rs.md) | | Configuration and deployment | [Store-RS deployment guides](../docs/source/deployment/store-rs/index.md) | From 0e626c572ba3ceee30cef4692bb981487bf65913 Mon Sep 17 00:00:00 2001 From: magicYang1573 <102229987+magicYang1573@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:31:43 +0800 Subject: [PATCH 07/14] [Store] Add session IDs to KV events for Dynamo session prefix indexer (#4683) --- .../api-reference/python/mooncake-store.md | 34 ++ mooncake-integration/store/store_py.cpp | 51 ++- .../store/store_py_internal.h | 3 +- .../include/kv_event/kv_event_publisher.h | 27 +- mooncake-store/include/metadata_store.h | 9 +- mooncake-store/include/object_metadata.h | 36 +++ mooncake-store/include/replica.h | 47 ++- mooncake-store/src/client_service.cpp | 35 +- .../src/kv_event/kv_event_publisher.cpp | 21 +- mooncake-store/src/master_service.cpp | 49 ++- mooncake-store/src/real_client.cpp | 12 + mooncake-store/src/rpc_service.cpp | 13 +- mooncake-store/tests/CMakeLists.txt | 3 +- .../tests/client_integration_test.cpp | 36 +++ mooncake-store/tests/file_storage_test.cpp | 14 +- .../ha/snapshot/batch_oplog/codec_test.cpp | 4 + .../tests/kv_event_publisher_test.cpp | 48 +++ mooncake-store/tests/session_tag_test.cpp | 301 ++++++++++++++++++ .../tests/test_session_kv_events.py | 259 +++++++++++++++ python/mooncake/structured_object_store.py | 17 +- .../store/test_structured_object_store.py | 16 + 21 files changed, 988 insertions(+), 47 deletions(-) create mode 100644 mooncake-store/tests/session_tag_test.cpp create mode 100644 mooncake-wheel/tests/test_session_kv_events.py diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md index 3c9eb87c43..2f3c743e89 100644 --- a/docs/source/api-reference/python/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -744,6 +744,40 @@ config = ReplicateConfig() config.prefer_alloc_in_same_node = True ``` +#### session_ids +**Type:** `Optional[List[str]]` +**Default:** `None` + +Optional per-key write attribution. Supply one entry for each input key (one +entry for a single-key write). Each ID must be valid UTF-8 and at most 256 bytes. +An empty string means no new attribution. The field does not change object identity, +replication, group membership, leases, or eviction policy. + +```python +config = ReplicateConfig() +config.session_ids = ["session-a", "session-a", "session-b"] +store.batch_put_from(keys, buffer_ptrs, buffer_sizes, config) +``` + +The tag becomes committed when the first replica of the write becomes readable. +A skipped duplicate Put leaves the original tag unchanged. Upsert with a +non-empty ID replaces attribution at commit; an omitted or empty ID preserves +the previous tag, including when storage is reallocated. Failed/preempted +writes do not commit their pending tags. Committed tags survive metadata +OpLog/snapshot recovery and standby promotion; old metadata defaults to no tag. + +With KV events enabled, each `stored` / `BlockStored` map carries a single +`session_id` string, or MessagePack `nil` when untagged. Replica availability +changes reuse the committed object tag. `removed` and `cleared` keep their +existing schemas: they retract availability, not a particular session's +association. This matches SGLang's single-session stored-event attribution; +it does not enumerate every session that has reused a shared object. + +Old RPC readers can ignore the compatible extension, but do not persist or +publish session tags. Upgrade the Store client and Master to enable attribution. +A storage node's disk scan alone cannot reconstruct metadata tags after all +Master metadata is lost; preserve the metadata recovery artifacts as well. + #### group_ids **Type:** `List[str] | None` **Default:** `None` diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 9080a079f2..895da7a8fd 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -1171,11 +1171,18 @@ class MooncakeStorePyWrapper { ReplicateConfig MakeIndexedConfig( const ReplicateConfig &config, const std::vector &original_indices) const { + ReplicateConfig indexed_config = config; + if (config.session_ids) { + indexed_config.session_ids->clear(); + for (size_t index : original_indices) { + indexed_config.session_ids->push_back( + config.session_ids->at(index)); + } + } if (!config.group_ids.has_value()) { - return config; + return indexed_config; } - ReplicateConfig indexed_config = config; std::vector group_ids; group_ids.reserve(original_indices.size()); for (size_t index : original_indices) { @@ -1188,11 +1195,20 @@ class MooncakeStorePyWrapper { ReplicateConfig MakeRepeatedIndexedConfig( const ReplicateConfig &config, const std::vector &original_indices, int repeat_count) const { + ReplicateConfig indexed_config = config; + if (config.session_ids) { + indexed_config.session_ids->clear(); + for (size_t index : original_indices) { + for (int i = 0; i < repeat_count; ++i) { + indexed_config.session_ids->push_back( + config.session_ids->at(index)); + } + } + } if (!config.group_ids.has_value()) { - return config; + return indexed_config; } - ReplicateConfig indexed_config = config; std::vector group_ids; group_ids.reserve(original_indices.size() * static_cast(repeat_count)); @@ -1208,6 +1224,11 @@ class MooncakeStorePyWrapper { std::vector ValidateGroupIdsForBatchConfig( const ReplicateConfig &config, size_t key_count, const char *operation_name) const { + if (!config.ValidSessionIds(key_count)) { + LOG(ERROR) << operation_name << ": invalid per-key session IDs"; + return std::vector(key_count, + to_py_ret(ErrorCode::INVALID_PARAMS)); + } if (config.group_ids.has_value() && config.group_ids->size() != key_count) { LOG(ERROR) << operation_name @@ -1222,6 +1243,8 @@ class MooncakeStorePyWrapper { const std::string &key, pybind11::object tensor, const ReplicateConfig &config = ReplicateConfig{}, int tp_rank = 0, int tp_size = 1, int split_dim = 0) { + if (!config.ValidSessionIds(1)) + return to_py_ret(ErrorCode::INVALID_PARAMS); try { std::vector tp_keys; tp_keys.reserve(tp_size); @@ -1234,7 +1257,13 @@ class MooncakeStorePyWrapper { if (!infos.has_value()) { return to_py_ret(ErrorCode::INVALID_PARAMS); } - auto results = batch_put_tensor_infos_impl(tp_keys, *infos, config); + ReplicateConfig tp_config = config; + if (config.session_ids) { + tp_config.session_ids = + std::vector(tp_size, config.SessionId()); + } + auto results = + batch_put_tensor_infos_impl(tp_keys, *infos, tp_config); for (int ret : results) { if (ret != 0) return ret; } @@ -2155,6 +2184,18 @@ PYBIND11_MODULE(_store, m) { &ReplicateConfig::prefer_alloc_in_same_node) .def_readwrite("data_type", &ReplicateConfig::data_type) .def_readwrite("group_ids", &ReplicateConfig::group_ids) + .def_property( + "session_ids", + [](const ReplicateConfig &config) + -> std::optional> { + if (!config.session_ids) return std::nullopt; + return *config.session_ids; + }, + [](ReplicateConfig &config, + std::optional> ids) { + config.session_ids.reset(); + if (ids) config.session_ids = std::move(*ids); + }) .def("__str__", [](const ReplicateConfig &config) { std::ostringstream oss; oss << config; diff --git a/mooncake-integration/store/store_py_internal.h b/mooncake-integration/store/store_py_internal.h index e43c3df105..83a3be052a 100644 --- a/mooncake-integration/store/store_py_internal.h +++ b/mooncake-integration/store/store_py_internal.h @@ -867,7 +867,8 @@ bool is_default_replicate_config(const ReplicateConfig &config) { !config.soft_pin_ttl_ms.has_value() && !config.with_hard_pin && config.preferred_segments.empty() && config.preferred_segment.empty() && - !config.prefer_alloc_in_same_node && !config.group_ids.has_value(); + !config.prefer_alloc_in_same_node && !config.group_ids.has_value() && + !config.session_ids.has_value(); } std::optional parse_parallel_axis_spec( diff --git a/mooncake-store/include/kv_event/kv_event_publisher.h b/mooncake-store/include/kv_event/kv_event_publisher.h index 9e72c03e13..1d26ff5d0d 100644 --- a/mooncake-store/include/kv_event/kv_event_publisher.h +++ b/mooncake-store/include/kv_event/kv_event_publisher.h @@ -34,11 +34,14 @@ class KvEventPublisher { // TenantId wrapper; the TenantId overloads preserve the existing API. void PublishStored(const std::string& object_key, const std::string& medium, const std::string& tenant_id = "", - const std::string& group_id = ""); + const std::string& group_id = "", + const std::string& session_id = ""); void PublishStored(const std::string& object_key, const std::string& medium, const TenantId& tenant_id, - const std::string& group_id = "") { - PublishStored(object_key, medium, tenant_id.value(), group_id); + const std::string& group_id = "", + const std::string& session_id = "") { + PublishStored(object_key, medium, tenant_id.value(), group_id, + session_id); } void PublishRemoved(const std::string& object_key, const std::string& medium, @@ -61,7 +64,8 @@ class KvEventPublisher { void PublishCommitted(const std::string& object_key, const std::vector& current_media, const std::string& tenant_id = "", - const std::string& group_id = ""); + const std::string& group_id = "", + const std::string& session_id = ""); // Synchronize replica availability after an internal metadata mutation. // previous_media is the medium set captured before the mutation; media that @@ -70,7 +74,8 @@ class KvEventPublisher { const std::vector& current_media, const std::string& tenant_id = "", const std::string& group_id = "", - const std::vector& previous_media = {}); + const std::vector& previous_media = {}, + const std::string& session_id = ""); // Publish removal of every medium the object had before deletion. void PublishObjectRemoved( @@ -95,6 +100,7 @@ class KvEventPublisher { std::string medium; std::string tenant_id; std::string group_id; + std::string session_id{}; }; void EnqueueBatch(std::vector events); @@ -132,19 +138,22 @@ class KvEventPublisher { bool enabled() const { return false; } void PublishStored(const std::string&, const std::string&, - const std::string& = "", const std::string& = "") {} - void PublishStored(const std::string&, const std::string&, const TenantId&, + const std::string& = "", const std::string& = "", const std::string& = "") {} + void PublishStored(const std::string&, const std::string&, const TenantId&, + const std::string& = "", const std::string& = "") {} void PublishRemoved(const std::string&, const std::string&, const std::string& = "", const std::string& = "") {} void PublishRemoved(const std::string&, const std::string&, const TenantId&, const std::string& = "") {} void PublishCleared(const std::string& = "") {} void PublishCommitted(const std::string&, const std::vector&, - const std::string& = "", const std::string& = "") {} + const std::string& = "", const std::string& = "", + const std::string& = "") {} void SyncObjectState(const std::string&, const std::vector&, const std::string& = "", const std::string& = "", - const std::vector& = {}) {} + const std::vector& = {}, + const std::string& = "") {} void PublishObjectRemoved(const std::string&, const std::string& = "", const std::string& = "", const std::vector& = {}) {} diff --git a/mooncake-store/include/metadata_store.h b/mooncake-store/include/metadata_store.h index 957d43a77b..08ff3c5598 100644 --- a/mooncake-store/include/metadata_store.h +++ b/mooncake-store/include/metadata_store.h @@ -39,10 +39,12 @@ struct StandbyObjectMetadata { ObjectDataType::UNKNOWN}; // Data type classification struct_pack::compatible hard_pinned; // Eviction protection + struct_pack::compatible session_id; + StandbyObjectMetadata() = default; YLT_REFL(StandbyObjectMetadata, client_id, size, replicas, group_id, - data_type, hard_pinned); + data_type, hard_pinned, session_id); // Check if this metadata has valid replicas bool HasReplicas() const { return !replicas.empty(); } @@ -106,8 +108,10 @@ struct MetadataPayload { struct_pack::compatible data_type; // Data type struct_pack::compatible hard_pinned; // Hard pin state + struct_pack::compatible session_id; + YLT_REFL(MetadataPayload, client_id, size, replicas, group_id, data_type, - hard_pinned); + hard_pinned, session_id); // Convert to StandbyObjectMetadata StandbyObjectMetadata ToStandbyMetadata() const { @@ -116,6 +120,7 @@ struct MetadataPayload { meta.size = size; meta.replicas = replicas; meta.group_id = group_id.value_or(""); + meta.session_id = session_id.value_or(""); meta.data_type = data_type.value_or(ObjectDataType::UNKNOWN); meta.hard_pinned = hard_pinned.value_or(false); return meta; diff --git a/mooncake-store/include/object_metadata.h b/mooncake-store/include/object_metadata.h index 0846327d13..046fbee130 100644 --- a/mooncake-store/include/object_metadata.h +++ b/mooncake-store/include/object_metadata.h @@ -106,6 +106,42 @@ struct ObjectMetadata { const TenantId tenant_id; const std::string user_key; + // Committed attribution only; protected by the owning metadata shard lock. + std::string session_id; + struct PendingSessionTag { + std::string session_id; + std::unordered_set eligible_replica_ids; + }; + std::optional pending_session_tag; + + void BeginSessionTag(std::string id, + const std::unordered_set& replicas) { + pending_session_tag.reset(); + if (!id.empty()) { + pending_session_tag = PendingSessionTag{std::move(id), replicas}; + } + } + + void CommitSessionTag(ReplicaID replica_id) { + if (pending_session_tag && + pending_session_tag->eligible_replica_ids.contains(replica_id)) { + session_id = std::move(pending_session_tag->session_id); + pending_session_tag.reset(); + } + } + + void ClearSessionTagIfNoViableReplica() { + if (pending_session_tag && !HasReplica([this](const Replica& replica) { + return pending_session_tag->eligible_replica_ids.contains( + replica.id()) && + replica.is_processing() && + !replica.has_invalid_mem_handle() && + !replica.has_invalid_nof_handle(); + })) { + pending_session_tag.reset(); + } + } + mutable SpinLock lock; // Authoritative lease: ungrouped objects own one; grouped objects share // the group's. Never null after construction. diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index eb1bcadafe..f35b82fbd4 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -3,6 +3,10 @@ #include #include +#include + +#include +#include #include #include @@ -125,13 +129,44 @@ struct ReplicateConfig { // always hash(tenant, key) and is decoupled from groups. std::optional> group_ids{}; - ReplicateConfig ForSingleKey(size_t key_index) const { - ReplicateConfig key_config = *this; - if (group_ids.has_value()) { - key_config.group_ids = - std::vector{group_ids->at(key_index)}; + // Per-key write attribution, independent of key identity and eviction + // groups. Compatible field keeps pre-session RPC payloads readable in both + // directions. + struct_pack::compatible, 1> session_ids; + + bool ValidSessionIds(size_t key_count) const { + return !session_ids || + (session_ids->size() == key_count && + std::all_of(session_ids->begin(), session_ids->end(), + [](const auto& id) { + if (id.size() > 256) return false; + auto it = id.begin(); + while (it != id.end()) { + namespace utf = boost::locale::utf; + if (!utf::is_valid_codepoint( + utf::utf_traits::decode( + it, id.end()))) { + return false; + } + } + return true; + })); + } + + std::string SessionId() const { + return session_ids && !session_ids->empty() ? session_ids->front() : ""; + } + + // Reuse a copy of the batch config; select IDs from the original batch. + void SetKeyIds(const ReplicateConfig& batch_config, size_t key_index) { + if (batch_config.group_ids.has_value()) { + group_ids = + std::vector{batch_config.group_ids->at(key_index)}; + } + if (batch_config.session_ids) { + session_ids = std::vector{ + batch_config.session_ids->at(key_index)}; } - return key_config; } friend std::ostream& operator<<(std::ostream& os, diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index b6f8a214d2..0a9473c0a7 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -2517,6 +2517,8 @@ void Client::healDanglingLocalDiskBatchStarts( std::vector retry_keys; std::vector> retry_slices; std::vector retry_positions; + ReplicateConfig retry_config = config; + if (retry_config.session_ids) retry_config.session_ids->clear(); for (size_t j = 0; j < evict_keys.size(); ++j) { if (!evicted[j]) { LOG(WARNING) << "Failed to evict dangling LOCAL_DISK replica for " @@ -2530,12 +2532,16 @@ void Client::healDanglingLocalDiskBatchStarts( retry_keys.emplace_back(keys[already_exists[evict_positions[j]]]); retry_slices.emplace_back( slice_lengths[already_exists[evict_positions[j]]]); + if (config.session_ids) { + retry_config.session_ids->push_back( + config.session_ids->at(already_exists[evict_positions[j]])); + } } if (retry_keys.empty()) { return; } auto retry_responses = - master_client_.BatchPutStart(retry_keys, retry_slices, config); + master_client_.BatchPutStart(retry_keys, retry_slices, retry_config); if (retry_responses.size() != retry_keys.size()) { return; } @@ -2568,6 +2574,8 @@ void Client::StartBatchPut(std::vector& ops, std::vector keys; std::vector> slice_lengths; std::vector active_indices; + ReplicateConfig active_config = config; + if (active_config.session_ids) active_config.session_ids->clear(); keys.reserve(ops.size()); slice_lengths.reserve(ops.size()); @@ -2587,6 +2595,9 @@ void Client::StartBatchPut(std::vector& ops, } active_indices.emplace_back(i); keys.emplace_back(op.key); + if (config.session_ids) { + active_config.session_ids->push_back(config.session_ids->at(i)); + } std::vector slice_sizes; slice_sizes.reserve(op.slices.size()); @@ -2601,7 +2612,7 @@ void Client::StartBatchPut(std::vector& ops, } auto start_responses = - master_client_.BatchPutStart(keys, slice_lengths, config); + master_client_.BatchPutStart(keys, slice_lengths, active_config); // Ensure response size matches request size if (start_responses.size() != active_indices.size()) { @@ -2651,7 +2662,7 @@ void Client::StartBatchPut(std::vector& ops, } healDanglingLocalDiskBatchStarts(ops, active_indices, already_exists, keys, - slice_lengths, config); + slice_lengths, active_config); for (size_t i : already_exists) { auto& op = ops[active_indices[i]]; if (!op.IsResolved() && op.replicas.empty()) { @@ -2667,6 +2678,8 @@ void Client::StartBatchUpsert(std::vector& ops, std::vector keys; std::vector> slice_lengths; std::vector active_indices; + ReplicateConfig active_config = config; + if (active_config.session_ids) active_config.session_ids->clear(); keys.reserve(ops.size()); slice_lengths.reserve(ops.size()); @@ -2686,6 +2699,9 @@ void Client::StartBatchUpsert(std::vector& ops, } active_indices.emplace_back(i); keys.emplace_back(op.key); + if (config.session_ids) { + active_config.session_ids->push_back(config.session_ids->at(i)); + } std::vector slice_sizes; slice_sizes.reserve(op.slices.size()); @@ -2700,7 +2716,7 @@ void Client::StartBatchUpsert(std::vector& ops, } auto start_responses = - master_client_.BatchUpsertStart(keys, slice_lengths, config); + master_client_.BatchUpsertStart(keys, slice_lengths, active_config); // Ensure response size matches request size if (start_responses.size() != active_indices.size()) { @@ -3581,6 +3597,10 @@ std::vector> Client::BatchPut( const std::vector& keys, std::vector>& batched_slices, const ReplicateConfig& config, const WriteBufferStager& stager) { + if (!config.ValidSessionIds(keys.size())) { + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; @@ -3627,6 +3647,10 @@ std::vector> Client::BatchUpsert( const std::vector& keys, std::vector>& batched_slices, const ReplicateConfig& config, const WriteBufferStager& stager) { + if (!config.ValidSessionIds(keys.size())) { + return std::vector>( + keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; @@ -3666,6 +3690,9 @@ Client::StartBatchPutForSizes(const std::vector& keys, const ReplicateConfig& config) { std::vector, ErrorCode>> results(keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + if (!config.ValidSessionIds(keys.size())) { + return results; + } if (keys.size() != object_sizes.size()) { LOG(ERROR) << "StartBatchPutForSizes size mismatch: keys=" << keys.size() << ", sizes=" << object_sizes.size(); diff --git a/mooncake-store/src/kv_event/kv_event_publisher.cpp b/mooncake-store/src/kv_event/kv_event_publisher.cpp index f680153a7c..47d20f8a08 100644 --- a/mooncake-store/src/kv_event/kv_event_publisher.cpp +++ b/mooncake-store/src/kv_event/kv_event_publisher.cpp @@ -54,7 +54,7 @@ size_t ComputeEventMapSize(bool is_stored, bool is_cleared, bool emit_legacy, map_size += 1; // object_key } if (is_stored) { - map_size += 3; // base_block_idx, parent_hash, token_ids + map_size += 4; // base_block_idx, parent_hash, token_ids, session_id if (emit_legacy) { map_size += 1; // parent_block_hash } @@ -174,7 +174,8 @@ KvEventPublisher::~KvEventPublisher() { void KvEventPublisher::PublishStored(const std::string& object_key, const std::string& medium, const std::string& tenant_id, - const std::string& group_id) { + const std::string& group_id, + const std::string& session_id) { if (!config_.enabled) { return; } @@ -185,7 +186,8 @@ void KvEventPublisher::PublishStored(const std::string& object_key, } std::vector events; events.push_back(PendingEvent{EventKind::kStored, object_key, medium, - NormalizeTenant(tenant_id), group_id}); + NormalizeTenant(tenant_id), group_id, + session_id}); EnqueueBatch(std::move(events)); } @@ -220,7 +222,7 @@ void KvEventPublisher::PublishCleared(const std::string& tenant_id) { void KvEventPublisher::PublishCommitted( const std::string& object_key, const std::vector& current_media, const std::string& tenant_id, - const std::string& group_id) { + const std::string& group_id, const std::string& session_id) { if (!config_.enabled) { return; } @@ -232,7 +234,7 @@ void KvEventPublisher::PublishCommitted( // the previous contents. for (const auto& medium : SortedMedia(NormalizeMedia(current_media))) { events.push_back(PendingEvent{EventKind::kStored, object_key, medium, - normalized_tenant, group_id}); + normalized_tenant, group_id, session_id}); } EnqueueBatch(std::move(events)); } @@ -256,8 +258,8 @@ void KvEventPublisher::PublishObjectRemoved( void KvEventPublisher::SyncObjectState( const std::string& object_key, const std::vector& current_media, const std::string& tenant_id, - const std::string& group_id, - const std::vector& previous_media) { + const std::string& group_id, const std::vector& previous_media, + const std::string& session_id) { if (!config_.enabled) { return; } @@ -275,7 +277,8 @@ void KvEventPublisher::SyncObjectState( for (const auto& medium : SortedMedia(new_set)) { if (!previous_set.contains(medium)) { events.push_back(PendingEvent{EventKind::kStored, object_key, - medium, normalized_tenant, group_id}); + medium, normalized_tenant, group_id, + session_id}); } } EnqueueBatch(std::move(events)); @@ -446,6 +449,8 @@ void KvEventPublisher::PublishBatch(const std::vector& batch) { } if (is_stored) { + packer.pack("session_id"); + PackOptionalString(packer, item.pending.session_id); packer.pack("base_block_idx"); packer.pack_nil(); packer.pack("parent_hash"); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index c89bdba2b4..8dce602e4c 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -3664,6 +3664,8 @@ tl::expected MasterService::InstallLegacyStandbyObjects( standby_meta.data_type, standby_meta.group_id, object.tenant_id, object.user_key)); (void)inserted; + it->second.metadata.session_id = + standby_meta.session_id.value_or(""); if (!standby_meta.group_id.empty()) { it->second.metadata.lease_ = RegisterGroupMember( object.tenant_id, object.user_key, standby_meta.group_id); @@ -4133,6 +4135,8 @@ tl::expected MasterService::RestoreFromStandbyState( standby_meta.data_type, standby_meta.group_id, object.tenant_id, object.user_key)); (void)inserted; + it->second.metadata.session_id = + standby_meta.session_id.value_or(""); if (!standby_meta.group_id.empty()) { it->second.metadata.lease_ = RegisterGroupMember(object.tenant_id, object.user_key, @@ -5383,6 +5387,8 @@ auto MasterService::InsertMetadata( return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } } + it->second.metadata.BeginSessionTag(config.SessionId(), + eligible_replica_ids); it->second.metadata.BeginSoftPinAction(soft_pin_request, std::move(eligible_replica_ids)); if (deadline_to_index) { @@ -5519,6 +5525,9 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, VLOG(1) << "key=" << key << ", value_length=" << slice_length << ", config=" << config << ", action=put_start_begin"; + if (!config.ValidSessionIds(1)) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } auto group_id_result = GetGroupIdForKey(config, 1, 0); if (!group_id_result) { return tl::make_unexpected(group_id_result.error()); @@ -5739,6 +5748,7 @@ auto MasterService::PutEnd(const UUID& client_id, const ObjectMeta& object_meta, metadata.PendingSoftPinOwnsReplica(replica.id())) { completed_pending_replica = true; } + metadata.CommitSessionTag(replica.id()); replica.mark_complete(); if (replica.is_dfs_replica() && dfs_allocator_) { const auto& desc = replica.get_dfs_descriptor(); @@ -6013,6 +6023,7 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, r.mark_removed(); }); metadata.ClearPendingSoftPinIfNoViableReplica(); + metadata.ClearSessionTagIfNoViableReplica(); tl::expected persist_result; if (remaining.empty()) { @@ -6045,6 +6056,7 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); EraseReplicasWithCacheTotalAccounting(metadata, target_pred); metadata.ClearPendingSoftPinIfNoViableReplica(); + metadata.ClearSessionTagIfNoViableReplica(); const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata); if (enable_multi_tenants_ && before_charge > after_charge) { auto release_result = metadata.quota_ledger.ReleaseCommitted( @@ -6205,6 +6217,9 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, VLOG(1) << "key=" << key << ", value_length=" << slice_length << ", config=" << config << ", action=upsert_start_begin"; + if (!config.ValidSessionIds(1)) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } auto group_id_result = GetGroupIdForKey(config, 1, 0); if (!group_id_result) { return tl::make_unexpected(group_id_result.error()); @@ -6216,6 +6231,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, bool dfs_allocation_failed = false; ReplicateConfig allocation_config = config; std::string allocation_group_id = group_id; + std::string allocation_session_id; std::optional allocation_committed_soft_pin_timeout; auto admit = @@ -6340,6 +6356,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, auto processing_replicas = metadata.PopReplicas(&Replica::fn_is_processing); metadata.ClearPendingSoftPinAction(); + metadata.pending_session_tag.reset(); if (!processing_replicas.empty()) { FreeDfsReplicas(key, processing_replicas); std::lock_guard lock(discarded_replicas_mutex_); @@ -6352,6 +6369,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, // If no COMPLETE replicas survive the preemption, this key // effectively does not exist β€” fall through to Case A. if (!metadata.HasReplica(&Replica::fn_is_completed)) { + allocation_session_id = metadata.session_id; allocation_committed_soft_pin_timeout = metadata.GetCommittedSoftPinTimeout(); if (allocation_committed_soft_pin_timeout && @@ -6376,12 +6394,17 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, // routed by hash(tenant, key); group_id does not affect routing. if (it == tenant_state.metadata.end()) { VLOG(1) << "key=" << key << ", action=upsert_start_case_a"; - return AllocateAndInsertMetadata( + auto result = AllocateAndInsertMetadata( shard, client_id, key, slice_length, allocation_config, writer_host_id, allocation_group_id, object_id.tenant_id, now, *soft_pin_request, allocation_committed_soft_pin_timeout, &dfs_allocation_failed); + if (result) { + tenant_state.metadata.at(key).metadata.session_id = + allocation_session_id; + } + return result; } else { // --- Step 2: key exists with COMPLETE replicas β†’ Case B or C // --- @@ -6464,6 +6487,8 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, eligible_replica_ids.insert(replica.id()); replica.mark_processing(); }); + metadata.BeginSessionTag(config.SessionId(), + eligible_replica_ids); metadata.BeginSoftPinAction( *soft_pin_request, std::move(eligible_replica_ids)); SyncCacheTotalAccounting(metadata); @@ -6507,6 +6532,7 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, } allocation_group_id = metadata.group_id; + allocation_session_id = metadata.session_id; const auto previous_kv_media = KvMediaSnapshot(metadata); TenantQuotaLedger replacement_charge; auto* quota_account = GetBoundTenantQuotaHandle(tenant_state); @@ -6614,6 +6640,9 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, return allocate_result; } auto new_it = tenant_state.metadata.find(key); + if (new_it != tenant_state.metadata.end()) { + new_it->second.metadata.session_id = allocation_session_id; + } if (has_replacement_charge) { if (new_it == tenant_state.metadata.end()) { auto rollback_result = @@ -6716,6 +6745,9 @@ MasterService::BatchUpsertStart(const UUID& client_id, tl::expected, ErrorCode>>( keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)); } + if (!config.ValidSessionIds(keys.size())) { + return {keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)}; + } if (config.group_ids.has_value() && config.group_ids->size() != keys.size()) { LOG(ERROR) << "BatchUpsertStart: group_ids.size()=" @@ -6728,8 +6760,9 @@ MasterService::BatchUpsertStart(const UUID& client_id, std::vector, ErrorCode>> results; results.reserve(keys.size()); + ReplicateConfig key_config = config; for (size_t i = 0; i < keys.size(); ++i) { - auto key_config = config.ForSingleKey(i); + key_config.SetKeyIds(config, i); results.emplace_back(UpsertStart(client_id, keys[i], tenant_id, slice_lengths[i], key_config)); } @@ -10427,6 +10460,7 @@ void MasterService::DiscardExpiredProcessingReplicas( if (!metadata.IsValid() || metadata.AllReplicas(&Replica::fn_is_completed)) { metadata.ClearPendingSoftPinIfNoViableReplica(); + metadata.ClearSessionTagIfNoViableReplica(); if (!metadata.IsValid()) { EraseMetadata(tenant_state, it, tenant_it->first, QuotaEraseMode::kFull, &shard); @@ -10494,6 +10528,7 @@ void MasterService::DiscardExpiredProcessingReplicas( auto replicas = metadata.PopReplicas(&Replica::fn_is_processing); metadata.ClearPendingSoftPinIfNoViableReplica(); + metadata.ClearSessionTagIfNoViableReplica(); if (!replicas.empty()) { FreeDfsReplicas(it->first, replicas); discarded_replicas.emplace_back(std::move(replicas), ttl); @@ -13380,7 +13415,8 @@ void MasterService::PublishKvStored(const std::string& key, return; } kv_event_publisher_->PublishCommitted(key, KvMediaForMetadata(metadata), - tenant_id.value(), metadata.group_id); + tenant_id.value(), metadata.group_id, + metadata.session_id); } void MasterService::SyncKvObjectState( @@ -13403,8 +13439,8 @@ void MasterService::SyncKvObjectState( return; } kv_event_publisher_->SyncObjectState(key, current_media, tenant_id.value(), - metadata.group_id, - previous_media_hint); + metadata.group_id, previous_media_hint, + metadata.session_id); } void MasterService::PublishKvRemoved( @@ -13605,6 +13641,7 @@ std::string MasterService::SerializeMetadataForOpLog( payload.client_id = metadata.client_id; payload.size = metadata.TotalLength(); payload.group_id = metadata.group_id; + payload.session_id = metadata.session_id; payload.data_type = metadata.data_type; payload.hard_pinned = metadata.IsHardPinned(); @@ -13631,6 +13668,7 @@ std::string MasterService::SerializeMetadataForOpLogWithoutMemReplicas( payload.client_id = metadata.client_id; payload.size = metadata.TotalLength(); payload.group_id = metadata.group_id; + payload.session_id = metadata.session_id; payload.data_type = metadata.data_type; payload.hard_pinned = metadata.IsHardPinned(); @@ -13655,6 +13693,7 @@ std::string MasterService::SerializeMetadataForOpLogFromReplicaDescriptors( payload.size = metadata.TotalLength(); payload.replicas = replicas; payload.group_id = metadata.group_id; + payload.session_id = metadata.session_id; payload.data_type = metadata.data_type; payload.hard_pinned = metadata.IsHardPinned(); auto result = struct_pack::serialize(payload); diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 94ae09d44a..937d38d9cb 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -7490,6 +7490,9 @@ std::vector RealClient::batch_put_session_start( return std::vector( keys.size(), static_cast(toInt(ErrorCode::INVALID_PARAMS))); } + if (!config.ValidSessionIds(keys.size())) { + return std::vector(keys.size(), toInt(ErrorCode::INVALID_PARAMS)); + } if (config.group_ids.has_value() && config.group_ids->size() != keys.size()) { LOG(ERROR) << "batch_put_session_start: group_ids.size()=" @@ -7541,6 +7544,15 @@ std::vector RealClient::batch_put_session_start( start_config.group_ids = std::move(filtered_group_ids); } + if (start_config.session_ids) { + std::vector ids; + ids.reserve(start_indices.size()); + for (size_t idx : start_indices) { + ids.push_back(start_config.session_ids->at(idx)); + } + start_config.session_ids = std::move(ids); + } + auto start_responses = client_->StartBatchPutForSizes(start_keys, start_sizes, start_config); std::vector keys_to_revoke; diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index de169b404e..dfced05598 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -519,6 +519,9 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, << " != slice_lengths.size()=" << slice_lengths.size(); results.assign(keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } else if (!config.ValidSessionIds(keys.size())) { + results.assign(keys.size(), + tl::make_unexpected(ErrorCode::INVALID_PARAMS)); } else if (config.group_ids.has_value() && config.group_ids->size() != keys.size()) { LOG(ERROR) << "BatchPutStart: group_ids.size()=" @@ -532,10 +535,10 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, } else if (config.prefer_alloc_in_same_node) { ReplicateConfig new_config = config; for (size_t i = 0; i < keys.size(); ++i) { - auto key_config = new_config.ForSingleKey(i); + new_config.SetKeyIds(config, i); auto result = master_service_.PutStart( client_id, keys[i], resolved_tenant_id.value(), - slice_lengths[i], key_config); + slice_lengths[i], new_config); results.emplace_back(result); if ((i == 0) && result.has_value()) { std::string preferred_segment; @@ -554,8 +557,9 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, } } } else { + ReplicateConfig key_config = config; for (size_t i = 0; i < keys.size(); ++i) { - auto key_config = config.ForSingleKey(i); + key_config.SetKeyIds(config, i); results.emplace_back(master_service_.PutStart( client_id, keys[i], resolved_tenant_id.value(), slice_lengths[i], key_config)); @@ -774,6 +778,9 @@ WrappedMasterService::BatchUpsertStart( << " != slice_lengths.size()=" << slice_lengths.size(); results.assign(keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } else if (!config.ValidSessionIds(keys.size())) { + results.assign(keys.size(), + tl::make_unexpected(ErrorCode::INVALID_PARAMS)); } else if (config.group_ids.has_value() && config.group_ids->size() != keys.size()) { LOG(ERROR) << "BatchUpsertStart: group_ids.size()=" diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 00fb1a2a7a..911c0db38e 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -68,6 +68,7 @@ set_tests_properties(replica_selection_env_opt_in_test add_store_test(eviction_strategy_test eviction_strategy_test.cpp) add_store_test(deadline_scheduler_test deadline_scheduler_test.cpp) add_store_test(background_worker_test background_worker_test.cpp) +add_store_test(session_tag_test session_tag_test.cpp) add_store_test(kv_event_publisher_test kv_event_publisher_test.cpp) if(ENABLE_KV_EVENTS) target_compile_definitions(kv_event_publisher_test @@ -156,7 +157,7 @@ add_test( NAME object_checksum_client_test COMMAND client_integration_test - "--gtest_filter=ObjectChecksumTest.*:ClientIntegrationTest.ObjectChecksumRejectsCorruptedObject:ClientIntegrationTest.BatchPutPreservesObjectChecksumPairing" + "--gtest_filter=ObjectChecksumTest.*:ClientIntegrationTest.ObjectChecksumRejectsCorruptedObject:ClientIntegrationTest.BatchPutPreservesObjectChecksumPairing:ClientIntegrationTest.TaggedBatchWritesSkipChecksumFailures" ) set_tests_properties(object_checksum_client_test PROPERTIES ENVIRONMENT "MOONCAKE_STORE_CHECKSUM=1") diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index dad8595511..d14f03c313 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -473,6 +473,42 @@ TEST_F(ClientIntegrationTest, BatchPutPreservesObjectChecksumPairing) { } } +TEST_F(ClientIntegrationTest, TaggedBatchWritesSkipChecksumFailures) { + if (!ClientObjectChecksumConfig::IsEnabledAtFirstUse()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + void* buffer = client_buffer_allocator_->allocate(4); + ASSERT_NE(buffer, nullptr); + std::memcpy(buffer, "data", 4); + for (bool upsert : {false, true}) { + const std::vector keys = { + "invalid-slice", upsert ? "tagged-upsert" : "tagged-put"}; + std::vector> slices{{Slice{nullptr, 1}}, + {Slice{buffer, 4}}}; + ReplicateConfig config; + config.session_ids = std::vector{"invalid", "valid"}; + auto results = upsert ? test_client_->BatchUpsert(keys, slices, config) + : test_client_->BatchPut(keys, slices, config); + ASSERT_EQ(results.size(), 2); + ASSERT_FALSE(results[0]); + EXPECT_EQ(results[0].error(), ErrorCode::INVALID_PARAMS); + ASSERT_TRUE(results[1]) << toString(results[1].error()); + auto query = test_client_->Query(keys[1]); + ASSERT_TRUE(query); + EXPECT_EQ(query->object_checksum, ComputeCrcChecksum("data", 4)); + + // A short tag vector must not become valid after filtering a key. + config.session_ids = std::vector{"wrong"}; + results = upsert ? test_client_->BatchUpsert(keys, slices, config) + : test_client_->BatchPut(keys, slices, config); + for (const auto& result : results) { + ASSERT_FALSE(result); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + } + } + client_buffer_allocator_->deallocate(buffer, 4); +} + // Test Remove operation TEST_F(ClientIntegrationTest, RemoveOperation) { const std::string test_data = "Test data for removal"; diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index 3f74941b16..fb73f2dad1 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -540,7 +540,8 @@ class FileStorageTest : public ::testing::Test { EXPECT_TRUE(unmount.has_value()); std::free(seg_ptr); } - void RunBatchPutHealWipeScenario(const std::string& data_path) { + void RunBatchPutHealWipeScenario(const std::string& data_path, + bool with_session_ids = false) { std::filesystem::path master_root = std::filesystem::path(data_path) / "heal_master_batch"; std::filesystem::create_directories(master_root); @@ -624,7 +625,7 @@ class FileStorageTest : public ::testing::Test { // One batch with both wiped keys and a fresh key. const std::string fresh_key = "heal_batch_fresh"; const std::string updated(512, 'y'); - std::vector keys = {wiped_keys[0], wiped_keys[1], fresh_key}; + std::vector keys = {fresh_key, wiped_keys[0], wiped_keys[1]}; std::vector> batched_slices; std::vector buffers; for (size_t i = 0; i < keys.size(); ++i) { @@ -636,6 +637,10 @@ class FileStorageTest : public ::testing::Test { } ReplicateConfig cfg; cfg.replica_num = 1; + if (with_session_ids) { + cfg.session_ids = + std::vector{"fresh", "wiped-a", "wiped-b"}; + } auto results = client.value()->BatchPut(keys, batched_slices, cfg); ASSERT_EQ(results.size(), keys.size()); for (size_t i = 0; i < results.size(); ++i) { @@ -1367,4 +1372,9 @@ TEST_F(FileStorageTest, RunBatchPutHealWipeScenario(data_path); } +TEST_F(FileStorageTest, + TaggedBatchPutAfterPhysicalWipeHealsDanglingLocalDiskReplica) { + RunBatchPutHealWipeScenario(data_path, true); +} + } // namespace mooncake diff --git a/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp b/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp index 43d3899294..2698883ed2 100644 --- a/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp +++ b/mooncake-store/tests/ha/snapshot/batch_oplog/codec_test.cpp @@ -64,6 +64,7 @@ TEST(BatchOpLogSnapshotCodecTest, ObjectChunkRoundTripsAndChecksEnvelope) { metadata.size = 8192; metadata.group_id = "group"; metadata.hard_pinned = true; + metadata.session_id = "session-snapshot"; std::vector objects{ {.tenant_id = "tenant", .key = "key", .metadata = metadata}}; auto encoded = EncodeBatchOpLogSnapshotObjectChunk(7, std::move(objects)); @@ -79,6 +80,8 @@ TEST(BatchOpLogSnapshotCodecTest, ObjectChunkRoundTripsAndChecksEnvelope) { EXPECT_EQ(8192u, decoded->objects[0].metadata.size); EXPECT_EQ("group", decoded->objects[0].metadata.group_id); EXPECT_TRUE(decoded->objects[0].metadata.hard_pinned.value_or(false)); + EXPECT_EQ(decoded->objects[0].metadata.session_id.value_or(""), + "session-snapshot"); EXPECT_FALSE(DecodeBatchOpLogSnapshotObjectChunk(encoded, 8, 1)); EXPECT_FALSE(DecodeBatchOpLogSnapshotObjectChunk(encoded, 7, 300000000)); @@ -94,6 +97,7 @@ TEST(BatchOpLogSnapshotCodecTest, OldObjectChunkDefaultsHardPinnedToFalse) { ASSERT_TRUE(decoded) << decoded.error(); EXPECT_FALSE(decoded->objects[0].metadata.hard_pinned.value_or(false)); + EXPECT_FALSE(decoded->objects[0].metadata.session_id.has_value()); } TEST(BatchOpLogSnapshotCodecTest, RejectsTruncatedArtifacts) { diff --git a/mooncake-store/tests/kv_event_publisher_test.cpp b/mooncake-store/tests/kv_event_publisher_test.cpp index e3dfe7f51f..55699105f3 100644 --- a/mooncake-store/tests/kv_event_publisher_test.cpp +++ b/mooncake-store/tests/kv_event_publisher_test.cpp @@ -23,6 +23,10 @@ TEST(KvEventPublisherTest, DisabledPublisherIsNoop) { KvEventPublisher publisher(config); EXPECT_FALSE(publisher.enabled()); publisher.PublishStored("42", "cpu"); + publisher.PublishStored("42", "cpu", TenantId::Default(), "", "session"); + publisher.PublishCommitted("42", {"cpu"}, "default", "", "session"); + publisher.SyncObjectState("42", {"disk"}, "default", "", {"cpu"}, + "session"); publisher.PublishRemoved("42", "cpu"); publisher.PublishCleared("tenant-a"); const auto stats = publisher.GetStats(); @@ -182,6 +186,50 @@ struct ScopedSubscriber { } // namespace +TEST(KvEventPublisherTest, SessionTagIsStoredAttributionOnly) { + for (bool legacy : {false, true}) { + auto endpoint = MakeIpcEndpoint(); + auto config = MakeEnabledConfig(endpoint); + config.emit_legacy_compat_fields = legacy; + KvEventPublisher publisher(config); + ScopedSubscriber subscriber; + ASSERT_TRUE(subscriber.Connect(endpoint)); + std::string session = "session-a"; + publisher.PublishStored("a", "cpu", "default", "group", session); + session = "changed-after-enqueue"; + publisher.PublishCommitted("b", {"cpu", "disk"}, "default", "", + "session-b"); + publisher.SyncObjectState("c", {"disk"}, "default", "", {"cpu"}, + "session-c"); + publisher.PublishStored("untagged", "cpu"); + publisher.PublishCleared(); + size_t received = 0; + while (received < 7) { + msgpack::object_handle payload; + ASSERT_TRUE(ReceivePayload(subscriber.sub, &payload)); + for (uint32_t i = 0; i < EventCount(payload); ++i) { + const auto& event = EventAt(payload, i); + const auto* tag = FindMapValue(event, "session_id"); + const auto kind = + MsgpackString(*FindMapValue(event, "event_type")); + if (kind == "stored") { + ASSERT_NE(tag, nullptr); + const auto key = + MsgpackString(*FindMapValue(event, "object_key")); + if (key == "untagged") + EXPECT_EQ(tag->type, msgpack::type::NIL); + else + EXPECT_EQ(MsgpackString(*tag), "session-" + key); + } else { + EXPECT_EQ(tag, nullptr); + } + ++received; + } + } + EXPECT_EQ(received, 7u); + } +} + TEST(KvEventPublisherTest, PublishesSglangObjectKeyOverZmq) { const std::string endpoint = MakeIpcEndpoint(); const std::string object_key = diff --git a/mooncake-store/tests/session_tag_test.cpp b/mooncake-store/tests/session_tag_test.cpp new file mode 100644 index 0000000000..c05c1bed41 --- /dev/null +++ b/mooncake-store/tests/session_tag_test.cpp @@ -0,0 +1,301 @@ +#include +#include + +#include "master_service_test_fixture.h" +#include "metadata_store.h" +#include "rpc_service.h" + +namespace { +// Count only allocations made by the config-splitting test on its own thread. +thread_local bool track_allocations = false; +thread_local size_t tracked_allocation_bytes = 0; +} // namespace + +void* operator new(size_t size) { + if (void* ptr = std::malloc(size ? size : 1)) { + if (track_allocations) tracked_allocation_bytes += size; + return ptr; + } + throw std::bad_alloc(); +} + +void operator delete(void* ptr) noexcept { std::free(ptr); } +void operator delete(void* ptr, size_t) noexcept { std::free(ptr); } + +namespace mooncake::test { +namespace { +// Pre-session RPC shape. Keep this fixture independent of ReplicateConfig. +struct LegacyReplicateConfig { + size_t replica_num{1}; + size_t nof_replica_num{0}; + size_t dfs_replica_num{0}; + SoftPinAction soft_pin_action{SoftPinAction::PRESERVE}; + std::optional soft_pin_ttl_ms; + bool with_hard_pin{false}; + std::vector preferred_segments; + std::string preferred_segment; + std::vector preferred_nof_segments; + bool prefer_alloc_in_same_node{false}; + ObjectDataType data_type{ObjectDataType::UNKNOWN}; + std::string host_id; + std::optional> group_ids; +}; + +struct LegacyMetadataPayload { + UUID client_id{0, 0}; + uint64_t size{0}; + std::vector replicas; + struct_pack::compatible group_id; + struct_pack::compatible data_type; + struct_pack::compatible hard_pinned; +}; + +std::string SessionTag(MasterService& service, const std::string& key) { + MasterServiceTestPeer::MetadataAccessorRO access( + &service, MasterServiceTestPeer(service).MakeObjectIdentityForRequest( + key, TenantId::Default())); + EXPECT_TRUE(access.Exists()); + return access.Exists() ? access.Get().session_id : ""; +} + +ReplicateConfig Tagged(std::string id) { + ReplicateConfig config; + config.session_ids = std::vector{std::move(id)}; + return config; +} +} // namespace + +TEST(SessionTagTest, RpcCompatibilityAndBatchValidation) { + EXPECT_EQ(struct_pack::get_type_code(), + struct_pack::get_type_code()); + LegacyReplicateConfig legacy; + legacy.group_ids = std::vector{"group"}; + auto decoded = struct_pack::deserialize( + struct_pack::serialize(legacy)); + ASSERT_TRUE(decoded); + EXPECT_FALSE(decoded->session_ids); + EXPECT_EQ(decoded->group_ids, legacy.group_ids); + + auto config = Tagged("session"); + config.group_ids = legacy.group_ids; + auto old_reader = struct_pack::deserialize( + struct_pack::serialize(config)); + ASSERT_TRUE(old_reader); + EXPECT_EQ(old_reader->group_ids, legacy.group_ids); + auto new_reader = struct_pack::deserialize( + struct_pack::serialize(config)); + ASSERT_TRUE(new_reader); + EXPECT_EQ(new_reader->SessionId(), "session"); + config.session_ids = std::vector{"a", "", "b"}; + config.group_ids = std::vector{"g1", "g2", "g3"}; + EXPECT_TRUE(config.ValidSessionIds(3)); + EXPECT_FALSE(config.ValidSessionIds(2)); + config.dfs_replica_num = 2; + config.soft_pin_action = SoftPinAction::ENABLE; + config.soft_pin_ttl_ms = 60000; + auto key_config = config; + key_config.preferred_segment = "selected-segment"; + for (size_t i = 0; i < 3; ++i) { + key_config.SetKeyIds(config, i); + ASSERT_EQ(key_config.session_ids->size(), 1); + EXPECT_EQ(key_config.SessionId(), config.session_ids->at(i)); + ASSERT_EQ(key_config.group_ids->size(), 1); + EXPECT_EQ(key_config.group_ids->front(), config.group_ids->at(i)); + EXPECT_EQ(key_config.preferred_segment, "selected-segment"); + EXPECT_EQ(key_config.dfs_replica_num, 2); + EXPECT_EQ(key_config.soft_pin_action, SoftPinAction::ENABLE); + EXPECT_EQ(key_config.soft_pin_ttl_ms, 60000); + } + EXPECT_EQ(config.session_ids->size(), 3); + EXPECT_EQ(config.group_ids->size(), 3); + EXPECT_TRUE(config.preferred_segment.empty()); + EXPECT_TRUE(Tagged(std::string(256, 'x')).ValidSessionIds(1)); + EXPECT_FALSE(Tagged(std::string(257, 'x')).ValidSessionIds(1)); +} + +TEST(SessionTagTest, BatchSessionIdSelectionAllocatesLinearly) { + size_t previous_bytes = 0; + for (size_t count : {1024, 4096}) { + ReplicateConfig config; + config.session_ids = + std::vector(count, std::string(64, 's')); + size_t selected_bytes = 0; + tracked_allocation_bytes = 0; + track_allocations = true; + { + auto key_config = config; + for (size_t i = 0; i < count; ++i) { + key_config.SetKeyIds(config, i); + selected_bytes += key_config.session_ids->front().size(); + } + } + track_allocations = false; + const size_t allocated_bytes = tracked_allocation_bytes; + EXPECT_EQ(selected_bytes, count * 64); + EXPECT_LT(allocated_bytes, count * 512); + if (previous_bytes) { + EXPECT_LE(allocated_bytes, previous_bytes * 5); + } + previous_bytes = allocated_bytes; + RecordProperty("allocated_bytes_" + std::to_string(count), + std::to_string(allocated_bytes)); + } +} + +TEST(SessionTagTest, SessionIdsRequireValidUtf8) { + for (const std::string id : {"", "ascii", "\xc2\x80", "会", + "\xf0\x9f\x98\x80", "\xf4\x8f\xbf\xbf"}) { + EXPECT_TRUE(Tagged(id).ValidSessionIds(1)); + } + for (const std::string id : {"\xff", "\x80", "\xc2", "\xe4\xb8", "\xc2x", + "\xc0\x80", "\xe0\x80\x80", "\xed\xa0\x80", + "\xf0\x80\x80\x80", "\xf4\x90\x80\x80"}) { + EXPECT_FALSE(Tagged(id).ValidSessionIds(1)); + } +} + +TEST(SessionTagTest, MetadataPayloadCompatibility) { + LegacyMetadataPayload legacy; + legacy.group_id = "group"; + legacy.hard_pinned = true; + auto decoded = struct_pack::deserialize( + struct_pack::serialize(legacy)); + ASSERT_TRUE(decoded); + EXPECT_FALSE(decoded->session_id); + EXPECT_EQ(decoded->group_id.value_or(""), "group"); + decoded->session_id = "session"; + auto old_reader = struct_pack::deserialize( + struct_pack::serialize(*decoded)); + ASSERT_TRUE(old_reader); + EXPECT_TRUE(old_reader->hard_pinned.value_or(false)); + EXPECT_EQ(old_reader->group_id.value_or(""), "group"); +} + +TEST_F(MasterServiceTest, SessionTagSurvivesAddingDiskReplica) { + MasterService service; + const auto client = PrepareSimpleSegment(service).client_id; + PutCompletedObject(service, client, "key", Tagged("writer")); + Replica disk(client, 1024, "local://session", ReplicaStatus::COMPLETE); + ASSERT_TRUE(service.AddReplica(client, "key", TenantId::Default(), disk)); + EXPECT_EQ(SessionTag(service, "key"), "writer"); +} + +TEST_F(MasterServiceTest, SessionTagCommitsAndSurvivesUpsertReplacement) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(0).build()); + PrepareSimpleSegment(service); + const auto client = generate_uuid(); + auto a = Tagged("a"); + ASSERT_TRUE(service.PutStart(client, "key", TenantId::Default(), 1024, a)); + EXPECT_EQ(SessionTag(service, "key"), ""); + ASSERT_TRUE(service.PutEnd(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "a"); + EXPECT_FALSE(service.PutStart(client, "key", TenantId::Default(), 1024, + Tagged("b"))); + EXPECT_EQ(SessionTag(service, "key"), "a"); + + // Same-size, absent tag; then size-changing, explicit tag. + ASSERT_TRUE( + service.UpsertStart(client, "key", TenantId::Default(), 1024, {})); + ASSERT_TRUE(service.UpsertEnd(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "a"); + ASSERT_TRUE(service.UpsertStart(client, "key", TenantId::Default(), 2048, + Tagged("b"))); + EXPECT_EQ(SessionTag(service, "key"), "a"); + EXPECT_FALSE(service.UpsertEnd(generate_uuid(), "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "a"); + ASSERT_TRUE(service.UpsertEnd(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "b"); + + ASSERT_TRUE(service.UpsertStart(client, "key", TenantId::Default(), 4096, + Tagged(""))); + ASSERT_TRUE(service.UpsertEnd(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "b"); + // Preempt an unfinished tagged write; its pending tag must not survive. + ASSERT_TRUE(service.UpsertStart(client, "key", TenantId::Default(), 4096, + Tagged("failed"))); + EXPECT_EQ(SessionTag(service, "key"), "b"); + ASSERT_TRUE( + service.UpsertStart(client, "key", TenantId::Default(), 4096, {})); + ASSERT_TRUE(service.UpsertEnd(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, "key"), "b"); + ASSERT_TRUE(service.Remove("key", TenantId::Default(), true)); + PutCompletedObject(service, client, "key", {}); + EXPECT_EQ(SessionTag(service, "key"), ""); +} + +TEST_F(MasterServiceTest, SessionTagValidationAndBatchUpsert) { + MasterService service; + PrepareSimpleSegment(service); + const auto client = generate_uuid(); + EXPECT_FALSE(service.PutStart(client, "bad", TenantId::Default(), 1024, + Tagged(std::string(257, 'x')))); + EXPECT_FALSE(service.ExistKey("bad", TenantId::Default()).value_or(false)); + auto config = Tagged("a"); + auto results = service.BatchUpsertStart( + client, {"a", "b"}, TenantId::Default(), {1024, 1024}, config); + ASSERT_EQ(results.size(), 2); + EXPECT_FALSE(results[0]); + EXPECT_FALSE(results[1]); + config.session_ids = std::vector{"a", "b"}; + results = service.BatchUpsertStart(client, {"a", "b"}, TenantId::Default(), + {1024, 1024}, config); + ASSERT_TRUE(results[0]); + ASSERT_TRUE(results[1]); + for (const auto& key : {"a", "b"}) { + ASSERT_TRUE(service.UpsertEnd(client, key, TenantId::Default(), + ReplicaType::MEMORY)); + EXPECT_EQ(SessionTag(service, key), key); + } +} + +TEST_F(MasterServiceTest, SessionTagRevokedWriteDoesNotTagRetry) { + MasterService service; + PrepareSimpleSegment(service); + const auto client = generate_uuid(); + ASSERT_TRUE(service.PutStart(client, "key", TenantId::Default(), 1024, + Tagged("revoked"))); + ASSERT_TRUE(service.PutRevoke(client, "key", TenantId::Default(), + ReplicaType::MEMORY)); + PutCompletedObject(service, client, "key", {}); + EXPECT_EQ(SessionTag(service, "key"), ""); +} + +TEST_F(MasterServiceTest, SessionTagOpLogAndStandbyRestore) { + MasterService service( + MasterServiceConfig::builder().set_default_kv_lease_ttl(0).build()); + PrepareSimpleSegment(service); + const auto client = generate_uuid(); + PutCompletedObject(service, client, "key", Tagged("committed")); + ASSERT_TRUE(service.UpsertStart(client, "key", TenantId::Default(), 1024, + Tagged("pending"))); + std::string payload; + { + MasterServiceTestPeer::MetadataAccessorRO access( + &service, + MasterServiceTestPeer(service).MakeObjectIdentityForRequest( + "key", TenantId::Default())); + payload = MasterServiceTestPeer::SerializeMetadataForOpLogForTesting( + service, access.Get()); + } + auto decoded = struct_pack::deserialize(payload); + ASSERT_TRUE(decoded); + EXPECT_EQ(decoded->session_id.value_or(""), "committed"); + auto restored_meta = decoded->ToStandbyMetadata(); + // Disk descriptor avoids coupling this test to an original allocator + // address. + Replica disk(generate_uuid(), 1024, "local://session", + ReplicaStatus::COMPLETE); + restored_meta.replicas = {disk.get_descriptor()}; + MasterService restored; + ASSERT_TRUE(restored.RestoreFromStandbySnapshot( + {{"default", "key", restored_meta}}, 0, {})); + EXPECT_EQ(SessionTag(restored, "key"), "committed"); +} +} // namespace mooncake::test diff --git a/mooncake-wheel/tests/test_session_kv_events.py b/mooncake-wheel/tests/test_session_kv_events.py new file mode 100644 index 0000000000..b86529a830 --- /dev/null +++ b/mooncake-wheel/tests/test_session_kv_events.py @@ -0,0 +1,259 @@ +"""Run with MOONCAKE_TEST_MASTER_BINARY pointing to an ENABLE_KV_EVENTS build.""" + +import ctypes +import os +import socket +import subprocess +import time + +import pytest + + +@pytest.fixture(params=[False, True], ids=["canonical", "legacy-fields"]) +def session_store(request, tmp_path): + binary = os.getenv("MOONCAKE_TEST_MASTER_BINARY") + if not binary: + pytest.skip("set MOONCAKE_TEST_MASTER_BINARY to run native KV event tests") + zmq = pytest.importorskip("zmq") + msgpack = pytest.importorskip("msgpack") + from mooncake.store import MooncakeDistributedStore + + # Reserve distinct local ports while selecting the endpoints. + reservations = [socket.socket() for _ in range(4)] + for reservation in reservations: + reservation.bind(("0.0.0.0", 0)) + rpc, metadata, events, admin = [s.getsockname()[1] for s in reservations] + for reservation in reservations: + reservation.close() + endpoint = f"tcp://127.0.0.1:{events}" + context = zmq.Context() + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt(zmq.SUBSCRIBE, b"") + subscriber.setsockopt(zmq.LINGER, 0) + subscriber.connect(endpoint) + store = MooncakeDistributedStore() + with (tmp_path / "master.log").open("w") as log: + master = subprocess.Popen( + [ + binary, + f"--rpc_port={rpc}", + f"--metrics_port={admin}", + "--enable_metric_reporting=false", + "--default_kv_lease_ttl=100ms", + "--enable_http_metadata_server=true", + f"--http_metadata_server_port={metadata}", + "--enable_kv_events=true", + f"--kv_events_bind_endpoint={endpoint}", + "--kv_events_backend_id=session-test", + f"--kv_events_emit_legacy_compat={str(request.param).lower()}", + ], + stdout=log, + stderr=subprocess.STDOUT, + ) + try: + deadline = time.monotonic() + 15 + while True: + assert master.poll() is None, (tmp_path / "master.log").read_text() + try: + with socket.create_connection(("127.0.0.1", rpc), timeout=0.1): + break + except OSError: + assert time.monotonic() < deadline, "Master startup timed out" + time.sleep(0.05) + assert ( + store.setup( + "127.0.0.1", + f"http://127.0.0.1:{metadata}/metadata", + 32 * 1024 * 1024, + 16 * 1024 * 1024, + "tcp", + "", + f"127.0.0.1:{rpc}", + ) + == 0 + ) + time.sleep(0.3) # Allow the PUB/SUB subscription to reach the publisher. + + def receive(count): + result = [] + while len(result) < count: + assert subscriber.poll(5000), "Missing KV event" + frames = subscriber.recv_multipart() + assert len(frames) == 3 + result.extend(msgpack.unpackb(frames[2], raw=False)[1]) + assert len(result) == count + return result + + yield store, receive, subscriber + finally: + store.close() + master.terminate() + try: + master.wait(timeout=10) + except subprocess.TimeoutExpired: + master.kill() + master.wait() + subscriber.close() + context.term() + + +def tagged(*ids): + from mooncake.store import ReplicateConfig + + config = ReplicateConfig() + config.session_ids = list(ids) + return config + + +def test_session_tags_over_rpc_and_zmq(session_store): + from mooncake.store import ReplicateConfig + + store, receive, subscriber = session_store + config = ReplicateConfig() + assert config.session_ids is None + config.session_ids = ["a"] + assert config.session_ids == ["a"] + config.session_ids = None + assert config.session_ids is None + + def expect_stored(key, session): + event = receive(1)[0] + assert event["event_type"] == "stored" + assert event["object_key"] == key + assert event["session_id"] == session + + assert store.put("shared", b"old", tagged("a")) == 0 + expect_stored("shared", "a") + # Duplicate Put and reads do not claim a new session association. + store.put("shared", b"new", tagged("b")) + assert store.get("shared") == b"old" + assert not subscriber.poll(300) + assert store.upsert("shared", b"new", tagged("b")) == 0 + changes = receive(2) + assert [e["session_id"] for e in changes if e["event_type"] == "stored"] == ["b"] + assert store.upsert("shared", b"same-size-and-larger", tagged("")) == 0 + changes = receive(2) + assert {event["event_type"] for event in changes} == {"stored", "removed"} + for event in changes: + if event["event_type"] == "stored": + assert event["session_id"] == "b" + else: + assert "session_id" not in event + assert store.upsert("shared", b"absent") == 0 + changes = receive(2) + assert [e["session_id"] for e in changes if e["event_type"] == "stored"] == ["b"] + + assert store.put_batch(["one", "two"], [b"1", b"2"], tagged("x", "")) == 0 + assert {e["object_key"]: e["session_id"] for e in receive(2)} == { + "one": "x", + "two": None, + } + assert store.upsert_batch(["one", "two"], [b"3", b"4"], tagged("y", "z")) == 0 + assert { + e["object_key"]: e["session_id"] + for e in receive(4) + if e["event_type"] == "stored" + } == { + "one": "y", + "two": "z", + } + + config = tagged("same-node-a", "same-node-b") + config.prefer_alloc_in_same_node = True + buffer = ctypes.create_string_buffer(b"ab") + pointer = ctypes.addressof(buffer) + assert store.register_buffer(pointer, 2) == 0 + try: + assert store.batch_put_from_multi_buffers( + ["local-a", "local-b"], [[pointer], [pointer + 1]], [[1], [1]], config + ) == [0, 0] + finally: + store.unregister_buffer(pointer) + assert (store.get("local-a"), store.get("local-b")) == (b"a", b"b") + assert {e["object_key"]: e["session_id"] for e in receive(2)} == { + "local-a": "same-node-a", + "local-b": "same-node-b", + } + + for invalid in (tagged("x"), tagged("x", "x" * 257), tagged("x", b"\xff")): + assert store.put_batch(["bad1", "bad2"], [b"1", b"2"], invalid) != 0 + assert store.is_exist("bad1") == 0 + for write in (store.put, store.upsert): + assert write("invalid-utf8", b"x", tagged(b"\xff")) != 0 + assert store.is_exist("invalid-utf8") == 0 + assert ( + store.upsert_batch(["bad1", "bad2"], [b"1", b"2"], tagged("valid", b"\xff")) + != 0 + ) + assert store.is_exist("bad1") == 0 + assert store.put("too-long", b"x", tagged("会" * 86)) != 0 + assert not subscriber.poll(300) + assert store.put("boundary", b"x", tagged("x" * 256)) == 0 + expect_stored("boundary", "x" * 256) + assert store.put("unicode-boundary", b"x", tagged("会" * 85 + "x")) == 0 + expect_stored("unicode-boundary", "会" * 85 + "x") + assert store.remove("shared", force=True) == 0 + assert "session_id" not in receive(1)[0] + assert store.put("shared", b"fresh") == 0 + expect_stored("shared", None) + + # An already-open write session is filtered before the second RPC batch. + assert store.batch_put_session_start(["held"], [4], tagged("original")) == [0] + results = store.batch_put_session_start( + ["held", "kept"], [4, 4], tagged("wrong", "kept-session") + ) + assert results[0] != 0 and results[1] == 0 + buffer = ctypes.create_string_buffer(b"data") + pointer = ctypes.addressof(buffer) + assert store.register_buffer(pointer, 4) == 0 + try: + assert store.batch_put_from_multi_buffer_ranges( + ["kept"], [[pointer]], [[4]], [[0]] + ) == [4] + assert store.batch_put_session_end(["kept"]) == [0] + expect_stored("kept", "kept-session") + finally: + store.unregister_buffer(pointer) + store.batch_put_session_revoke(["held"]) + assert not subscriber.poll(300) + + +def test_tensor_filtering_and_tp_session_mapping(session_store): + torch = pytest.importorskip("torch") + store, receive, subscriber = session_store + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) + results = store.batch_pub_tensor( + ["invalid", "valid"], [None, tensor], tagged("wrong", "right") + ) + assert results[0] != 0 and results[1] == 0 + assert receive(1)[0]["session_id"] == "right" + config = tagged("tp-session") + config.group_ids = ["tp-left", "tp-right"] + assert store.pub_tensor_with_tp("tp", tensor, config, tp_size=2) == 0 + events = receive(2) + assert [e["session_id"] for e in events] == ["tp-session"] * 2 + assert sorted(e["group_id"] for e in events) == ["tp-left", "tp-right"] + results = store.batch_pub_tensor_with_tp( + ["tp-a", "tp-b"], [tensor, tensor], tagged("a", "b"), tp_size=2 + ) + assert results == [0, 0] + assert sorted(e["session_id"] for e in receive(4)) == ["a", "a", "b", "b"] + assert not subscriber.poll(300) + + +def test_structured_session_expansion_preserves_native_write_options(): + from mooncake.store import ReplicateConfig, SoftPinAction + from mooncake.structured_object_store import _config_for_grouped_keys + + config = ReplicateConfig() + config.dfs_replica_num = 1 + config.soft_pin_action = SoftPinAction.ENABLE + config.soft_pin_ttl_ms = 60000 + config.session_ids = ["session"] + expanded = _config_for_grouped_keys(config, 2) + assert expanded.session_ids == ["session", "session"] + assert expanded.dfs_replica_num == 1 + assert expanded.soft_pin_action == SoftPinAction.ENABLE + assert expanded.soft_pin_ttl_ms == 60000 + assert config.session_ids == ["session"] + assert expanded.group_ids is None diff --git a/python/mooncake/structured_object_store.py b/python/mooncake/structured_object_store.py index aab6ad99aa..db902987d6 100644 --- a/python/mooncake/structured_object_store.py +++ b/python/mooncake/structured_object_store.py @@ -6737,7 +6737,9 @@ def _call_write_with_optional_config(fn: Any, *args: Any, config: Any = None) -> _WRITE_CONFIG_COPY_FIELDS = ( "data_type", + "dfs_replica_num", "group_ids", + "session_ids", "nof_replica_num", "prefer_alloc_in_same_node", "preferred_nof_segments", @@ -6745,7 +6747,8 @@ def _call_write_with_optional_config(fn: Any, *args: Any, config: Any = None) -> "preferred_segments", "replica_num", "with_hard_pin", - "with_soft_pin", + "soft_pin_action", + "soft_pin_ttl_ms", ) @@ -6753,6 +6756,18 @@ def _config_for_grouped_keys(config: Any, key_count: int) -> Any: """Return a write config whose group_ids match the physical key count.""" if config is None or key_count <= 0: return config + session_ids = getattr(config, "session_ids", None) + if session_ids is not None: + session_ids = ( + [session_ids] if isinstance(session_ids, str) else list(session_ids) + ) + if len(session_ids) != 1: + raise ValueError( + "config.session_ids must contain exactly one logical session id" + ) + if key_count != 1 or isinstance(config.session_ids, str): + config = _copy_write_config(config) + config.session_ids = session_ids * key_count group_ids = getattr(config, "group_ids", None) if group_ids is None: return config diff --git a/python/tests/store/test_structured_object_store.py b/python/tests/store/test_structured_object_store.py index 40295898dd..cf807c02c2 100644 --- a/python/tests/store/test_structured_object_store.py +++ b/python/tests/store/test_structured_object_store.py @@ -5143,3 +5143,19 @@ def test_dataproto_helper_recursive_manifest_export_append_and_overwrite() -> No assert ref.encoded_non_tensor["samples"]["codec"] == "utf8_ragged" assert overwritten["non_tensor_batch"]["samples"].tolist() == ["a", "b", "c", "d"] assert old_manifest_key not in store.objects + + +def test_session_config_expands_with_physical_keys_without_mutating_input(): + from types import SimpleNamespace + from mooncake.structured_object_store import _config_for_grouped_keys + + config = SimpleNamespace(group_ids=["group"], session_ids=["session"]) + expanded = _config_for_grouped_keys(config, 3) + assert expanded.group_ids == ["group"] * 3 + assert expanded.session_ids == ["session"] * 3 + assert config.group_ids == ["group"] + assert config.session_ids == ["session"] + session_only = SimpleNamespace(session_ids=["session"]) + assert _config_for_grouped_keys(session_only, 2).session_ids == ["session"] * 2 + with pytest.raises(ValueError, match="session_ids"): + _config_for_grouped_keys(SimpleNamespace(session_ids=["a", "b"]), 2) From d782a5998f6dc3395f2e0b6487e09229a7e4e1e6 Mon Sep 17 00:00:00 2001 From: Daegyu Han <30819476+daegyu94@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:11:13 +0900 Subject: [PATCH 08/14] [Store] Reuse validated shard paths on cached DFS I/O (#4755) Co-authored-by: Daegyu Han --- .../distributed_storage_backend.cpp | 24 ++++++++++--- mooncake-store/tests/dfs_posix_test.cpp | 36 +++++++++++++++++++ 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp index 68549fd188..f436638874 100644 --- a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp +++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp @@ -216,6 +216,24 @@ tl::expected DistributedStorageBackend::Init() { tl::expected DistributedStorageBackend::GetOrOpenShard( const DistributedFSDescriptor& descriptor) { + if (descriptor.shard_idx < 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + ShardFile* shard = nullptr; + { + std::lock_guard cache_lock(shard_files_mutex_); + auto existing = shard_files_.find(descriptor.shard_idx); + if (existing != shard_files_.end()) shard = existing->second.get(); + } + if (shard) { + std::lock_guard shard_lock(shard->mutex); + // Only the exact path validated when this fd was opened can bypass + // normalization. Aliases and changed descriptors take the full + // validation path below. Cache entries live until backend teardown. + if (shard->fd >= 0 && shard->path == descriptor.file_path) return shard; + } + const auto path = std::filesystem::path(descriptor.file_path).lexically_normal(); auto root = std::filesystem::path(root_dir_).lexically_normal(); @@ -223,13 +241,11 @@ DistributedStorageBackend::GetOrOpenShard( auto parent = path.parent_path(); if (parent.empty()) parent = "."; const auto index = ParseShardFileName(path.filename().string()); - if (descriptor.shard_idx < 0 || !index || *index != descriptor.shard_idx || - parent != root) { + if (!index || *index != descriptor.shard_idx || parent != root) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - ShardFile* shard; - { + if (!shard) { std::lock_guard cache_lock(shard_files_mutex_); auto existing = shard_files_.find(descriptor.shard_idx); if (existing == shard_files_.end()) { diff --git a/mooncake-store/tests/dfs_posix_test.cpp b/mooncake-store/tests/dfs_posix_test.cpp index 8a9b939b0a..1c32cb2db0 100644 --- a/mooncake-store/tests/dfs_posix_test.cpp +++ b/mooncake-store/tests/dfs_posix_test.cpp @@ -1425,6 +1425,42 @@ TEST_F(DfsBackendTest, LazyOpenRejectsMissingAndMalformedShardFiles) { EXPECT_EQ(write[0].error(), ErrorCode::INVALID_PARAMS); } +// Once a shard is open, only the exact path validated at open time skips +// normalization. Aliases and changed descriptors must still be validated +// against the cached shard without reopening it. +TEST_F(DfsBackendTest, CachedShardRevalidatesChangedDescriptors) { + ControlledPosixFsAdapter* controlled = nullptr; + auto backend = MakeControlledBackend(tmp_->path(), &controlled); + ASSERT_TRUE(backend->Init().has_value()); + + AlignedBuffer buf(4096); + ASSERT_NE(buf.data(), nullptr); + const auto read = [&](const std::string& path, int shard_idx) { + auto results = backend->BatchRead({{"cached", + {path, 0, 4096, 4096, shard_idx}, + {{buf.data(), buf.size()}}}}); + return results.at(0); + }; + const std::filesystem::path root(tmp_->path()); + const auto name = std::filesystem::path(ShardPath(0)).filename(); + + ASSERT_TRUE(read(ShardPath(0), 0)); + ASSERT_TRUE(read(ShardPath(0), 0)); + EXPECT_TRUE(read((root / "." / name).string(), 0)); + + for (const auto& [path, shard_idx] : + std::vector>{ + {ShardPath(1), 0}, + {ShardPath(0), 1}, + {(root / ".." / name).string(), 0}}) { + SCOPED_TRACE(path + " shard_idx=" + std::to_string(shard_idx)); + auto result = read(path, shard_idx); + ASSERT_FALSE(result); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + } + EXPECT_EQ(controlled->OpenCallCount(), 1); +} + TEST_F(DfsBackendTest, ConcurrentLazyOpenPreservesIndependentWrites) { PosixFsAdapter adapter; ASSERT_TRUE(adapter.Init(tmp_->path())); From c6f1d08e1cb8f20b927c57009a0ad864a8db209a Mon Sep 17 00:00:00 2001 From: 2xdevv Date: Mon, 19 Oct 2026 15:49:40 +0900 Subject: [PATCH 09/14] [Store] Allow DFS-only replica placement (#4637) Co-authored-by: 2xdevv <2xdevv@gmail.com> --- .../api-reference/cpp/mooncake-store.md | 2 +- .../api-reference/python/mooncake-store.md | 7 +- .../mooncake-store-deployment-guide.md | 13 +- mooncake-store/src/master_service.cpp | 14 +- .../dfs_immutable_bucket_client_test.cpp | 54 +++++++ mooncake-store/tests/dfs_sync_client_test.cpp | 122 +++++++++++++++ mooncake-store/tests/master_service_test.cpp | 144 ++++++++++++++++++ 7 files changed, 338 insertions(+), 18 deletions(-) diff --git a/docs/source/api-reference/cpp/mooncake-store.md b/docs/source/api-reference/cpp/mooncake-store.md index 49d9ca4ed8..a2bb65cf4d 100644 --- a/docs/source/api-reference/cpp/mooncake-store.md +++ b/docs/source/api-reference/cpp/mooncake-store.md @@ -64,7 +64,7 @@ struct ReplicateConfig { }; ``` -`dfs_replica_num` may currently be `0` or `1`. When it is `1`, `replica_num >= 1` is required, so DFS-only placement is not supported. DFS replicas currently support only the `default` tenant and require the master and client DFS backends to be configured with the same shared root and shard layout. See the {ref}`DFS deployment documentation ` for setup and lifecycle limitations. Native C++ clients must initialize a `DistributedStorageBackend` and attach it with `SetDfsStorageBackend()` before issuing DFS reads or writes. DFS descriptors are carried by each `PutStart`, `UpsertStart`, or query response; there is no client-side descriptor cache. The Python/RealClient setup path attaches the backend through `FileStorage`. +`dfs_replica_num` may currently be `0` or `1`. DFS-only placement (`replica_num = 0`, `nof_replica_num = 0`, `dfs_replica_num = 1`) is supported; the DFS replica is then the object's only copy, so DFS eviction or removal deletes the object. DFS replicas currently support only the `default` tenant and require the master and client DFS backends to be configured with the same shared root and shard layout. See the {ref}`DFS deployment documentation ` for setup and lifecycle limitations. Native C++ clients must initialize a `DistributedStorageBackend` and attach it with `SetDfsStorageBackend()` before issuing DFS reads or writes. DFS descriptors are carried by each `PutStart`, `UpsertStart`, or query response; there is no client-side descriptor cache. The Python/RealClient setup path attaches the backend through `FileStorage`. Soft pinning starts when the first replica becomes readable and has a fixed lifetime: reads do not extend it. `PRESERVE` keeps the committed deadline on an diff --git a/docs/source/api-reference/python/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md index 2f3c743e89..2024ca1dca 100644 --- a/docs/source/api-reference/python/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -667,9 +667,10 @@ config.nof_replica_num = 1 **Default:** `0` **Status:** **Work in progress; development and evaluation only.** **Description:** Requests an additional replica in the configured shared -distributed filesystem. The supported values are currently `0` and `1`. When -set to `1`, `replica_num` must be at least `1`, so DFS-only placement is not -supported. DFS replicas currently support only the `default` tenant. +distributed filesystem. The supported values are currently `0` and `1`. +DFS-only placement (`replica_num = 0` with `dfs_replica_num = 1`) is supported; +the DFS replica is then the object's only copy, so DFS eviction or removal +deletes the object. DFS replicas currently support only the `default` tenant. ```python config = ReplicateConfig() diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 004d50d5eb..d185ce7790 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -998,9 +998,9 @@ config.dfs_replica_num = 1 store.put("key", b"value", config) ``` -`dfs_replica_num` may currently be `0` or `1`. A DFS replica must be requested -with at least one memory replica (`replica_num >= 1`), so DFS-only placement is -not supported. +`dfs_replica_num` may currently be `0` or `1`. DFS-only placement is supported +by setting `replica_num = 0` and `nof_replica_num = 0`; the object then has no +memory replica, and DFS eviction or removal deletes it. Allocation first tries the key's hash-selected DFS shard, then tries other ready shards if that shard has no suitable extent. This lets new shards accept @@ -1022,7 +1022,9 @@ failed key is revoked without downgrading successful keys. For a same-size `Upsert`, if either the existing object or the new request has a DFS replica, the requested memory, NoF, and DFS replica counts must match the existing topology. A different-size update releases the old placement and -allocates a new topology. +allocates a new topology. For example, after the memory replica of a +memory-plus-DFS object is evicted, a same-size `Upsert` must request the +DFS-only topology. On reads, the master returns the readable replica list through the normal query path, and the client selects the first complete replica. If it selects DFS, any @@ -1032,8 +1034,7 @@ reads for that descriptor. #### Current limitations - Only the `default` tenant is supported. -- `dfs_replica_num` must be `0` or `1`, and `replica_num >= 1` is required when - it is enabled. +- `dfs_replica_num` must be `0` or `1`. - C and Rust clients cannot currently request or access descriptor-based DFS: their replication configuration does not expose `dfs_replica_num`, and their setup API cannot initialize the distributed `FileStorage` backend. Use the diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 8dce602e4c..6deb936492 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -5466,11 +5466,10 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, << ", key_size=" << key.size() << ", error=invalid_params"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - if (config.dfs_replica_num > 1 || - (config.dfs_replica_num > 0 && config.replica_num == 0)) { - LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num + if (config.dfs_replica_num > 1) { + LOG(ERROR) << "key=" << key << ", dfs_replica_num=" << config.dfs_replica_num - << ", error=invalid_dfs_replica_config"; + << ", error=dfs_replica_num_must_be_0_or_1"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } if (config.dfs_replica_num > 0 && !object_id.tenant_id.IsDefault()) { @@ -6155,11 +6154,10 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, << ", key_size=" << key.size() << ", error=invalid_params"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - if (config.dfs_replica_num > 1 || - (config.dfs_replica_num > 0 && config.replica_num == 0)) { - LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num + if (config.dfs_replica_num > 1) { + LOG(ERROR) << "key=" << key << ", dfs_replica_num=" << config.dfs_replica_num - << ", error=invalid_dfs_replica_config"; + << ", error=dfs_replica_num_must_be_0_or_1"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } if (config.dfs_replica_num > 0 && !object_id.tenant_id.IsDefault()) { diff --git a/mooncake-store/tests/dfs_immutable_bucket_client_test.cpp b/mooncake-store/tests/dfs_immutable_bucket_client_test.cpp index fd031b053d..22262b112d 100644 --- a/mooncake-store/tests/dfs_immutable_bucket_client_test.cpp +++ b/mooncake-store/tests/dfs_immutable_bucket_client_test.cpp @@ -242,6 +242,60 @@ TEST_F(DfsImmutableBucketClientTest, ExhaustionEvictsWholeBucketAndRetries) { EXPECT_EQ(third_query->replicas.size(), 2u); } +TEST_F(DfsImmutableBucketClientTest, DfsOnlyObjectsLiveAndDieWithTheirBucket) { + ReplicateConfig dfs_only; + dfs_only.replica_num = 0; + dfs_only.dfs_replica_num = 1; + auto expect_value = [this](const std::string& key, + const std::string& expected) { + auto query = writer_->Query(key); + ASSERT_TRUE(query) << key; + ASSERT_EQ(query->replicas.size(), 1u) << key; + ASSERT_TRUE(query->replicas[0].is_dfs_replica()) << key; + // The fixture's zero lease TTL makes Client::Get report LEASE_EXPIRED, + // so read the descriptor through the backend like the other tests. + std::string output(expected.size(), '\0'); + auto reads = + backend_->BatchRead({{key, + query->replicas[0].get_dfs_descriptor(), + {{output.data(), output.size()}}}}); + ASSERT_TRUE(reads[0]) << key; + EXPECT_EQ(output, expected) << key; + }; + + std::string first(4096, '1'); + std::string second(4096, '2'); + std::string third(4096, '3'); + auto first_slices = Slices(first); + auto second_slices = Slices(second); + auto third_slices = Slices(third); + ASSERT_TRUE(writer_->Put("first", first_slices, dfs_only)); + ASSERT_TRUE(writer_->Put("second", second_slices, dfs_only)); + // Both buckets are full: the third write forces a whole-bucket eviction. + // With no memory replica, the evicted object disappears entirely. + ASSERT_TRUE(writer_->Put("third", third_slices, dfs_only)); + auto first_query = writer_->Query("first"); + ASSERT_FALSE(first_query); + EXPECT_EQ(first_query.error(), ErrorCode::OBJECT_NOT_FOUND); + expect_value("third", third); + + adapter_->FailNextWrite(); + std::string failed(4096, 'F'); + auto failed_slices = Slices(failed); + auto failed_put = writer_->Put("failed", failed_slices, dfs_only); + ASSERT_FALSE(failed_put); + EXPECT_EQ(failed_put.error(), ErrorCode::FILE_WRITE_FAIL); + auto missing = writer_->Query("failed"); + ASSERT_FALSE(missing); + EXPECT_EQ(missing.error(), ErrorCode::OBJECT_NOT_FOUND); + + // Same-size upsert reallocates into a new bucket. + std::string replacement(4096, 'R'); + auto replacement_slices = Slices(replacement); + ASSERT_TRUE(writer_->Upsert("third", replacement_slices, dfs_only)); + expect_value("third", replacement); +} + TEST_F(DfsImmutableBucketLruClientTest, BatchQueryRefreshesDfsBucketLru) { std::string first(4096, '1'); std::string second(4096, '2'); diff --git a/mooncake-store/tests/dfs_sync_client_test.cpp b/mooncake-store/tests/dfs_sync_client_test.cpp index 7505d1cca5..119b1b9b6d 100644 --- a/mooncake-store/tests/dfs_sync_client_test.cpp +++ b/mooncake-store/tests/dfs_sync_client_test.cpp @@ -133,6 +133,34 @@ class DfsSyncClientTest : public ::testing::Test { return config; } + static ReplicateConfig DfsOnlyConfig() { + ReplicateConfig config; + config.replica_num = 0; + config.dfs_replica_num = 1; + return config; + } + + // Reads through Query + Client::Get and checks that DFS is the only + // replica. + void ExpectDfsOnlyValue(const std::string& key, + const std::string& expected) { + auto query = writer_->Query(key); + ASSERT_TRUE(query.has_value()) << key; + last_read_lease_deadline_ = query->lease_timeout; + ASSERT_EQ(query->replicas.size(), 1) << key; + ASSERT_TRUE(query->replicas[0].is_dfs_replica()) << key; + std::string output(expected.size(), '\0'); + std::vector slices{{output.data(), output.size()}}; + ASSERT_TRUE(writer_->Get(key, *query, slices).has_value()) << key; + EXPECT_EQ(output, expected) << key; + } + + void ExpectNotFound(const std::string& key) { + auto query = writer_->Query(key); + ASSERT_FALSE(query.has_value()) << key; + EXPECT_EQ(query.error(), ErrorCode::OBJECT_NOT_FOUND) << key; + } + tl::expected QueryDfsOnly(const std::string& key) { auto query = writer_->Query(key); if (!query) { @@ -570,4 +598,98 @@ TEST_F(DfsSyncClientTest, DfsMetricsAppearInSerializedOutput) { EXPECT_NE(summary->find("=== DFS Metrics Summary ==="), std::string::npos); } +TEST_F(DfsSyncClientTest, DfsOnlyPutGetRemoveAndReput) { + const std::string key = "dfs_only_put"; + std::string value(4096, 'a'); + std::vector slices{{value.data(), value.size()}}; + ASSERT_TRUE(writer_->Put(key, slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, value); + + WaitForReadLeaseExpiry(); + ASSERT_TRUE(writer_->Remove(key).has_value()); + ExpectNotFound(key); + + std::string larger(8192, 'b'); + std::vector larger_slices{{larger.data(), larger.size()}}; + ASSERT_TRUE(writer_->Put(key, larger_slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, larger); +} + +TEST_F(DfsSyncClientTest, DfsOnlyBatchFailureOnlyRevokesFailedKey) { + std::vector keys{"dfs_only_ok", "dfs_only_fail"}; + std::vector values{std::string(4096, 'c'), + std::string(4096, 'd')}; + auto slices = MakeSlices(values); + adapter_->FailWriteCall(adapter_->WriteCalls() + 2); + + auto results = writer_->BatchPut(keys, slices, DfsOnlyConfig()); + ASSERT_EQ(results.size(), keys.size()); + EXPECT_TRUE(results[0].has_value()); + ASSERT_FALSE(results[1].has_value()); + EXPECT_EQ(results[1].error(), ErrorCode::FILE_WRITE_FAIL); + ExpectDfsOnlyValue(keys[0], values[0]); + ExpectNotFound(keys[1]); + + // The revoked key and its extent can be reused. + std::vector retry{{values[1].data(), values[1].size()}}; + ASSERT_TRUE(writer_->Put(keys[1], retry, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(keys[1], values[1]); +} + +TEST_F(DfsSyncClientTest, DfsOnlyUpsertSameSizeResizeAndBatch) { + const std::string key = "dfs_only_upsert"; + std::string initial(4096, 'e'); + std::vector initial_slices{{initial.data(), initial.size()}}; + ASSERT_TRUE( + writer_->Upsert(key, initial_slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, initial); + + WaitForReadLeaseExpiry(); + std::string same_size(4096, 'f'); + std::vector same_slices{{same_size.data(), same_size.size()}}; + ASSERT_TRUE(writer_->Upsert(key, same_slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, same_size); + + WaitForReadLeaseExpiry(); + std::string resized(12288, 'g'); + std::vector resized_slices{{resized.data(), resized.size()}}; + ASSERT_TRUE( + writer_->Upsert(key, resized_slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, resized); + + WaitForReadLeaseExpiry(); + std::vector keys{key, "dfs_only_upsert_new"}; + std::vector values{std::string(4096, 'h'), + std::string(4096, 'i')}; + auto batch_slices = MakeSlices(values); + auto results = writer_->BatchUpsert(keys, batch_slices, DfsOnlyConfig()); + ASSERT_EQ(results.size(), keys.size()); + ASSERT_TRUE(results[0].has_value()); + ASSERT_TRUE(results[1].has_value()); + ExpectDfsOnlyValue(keys[0], values[0]); + ExpectDfsOnlyValue(keys[1], values[1]); +} + +TEST_F(DfsSyncClientTest, DfsOnlyUpsertMatchesMemoryLostTopology) { + const std::string key = "dfs_only_memory_lost"; + std::string initial(4096, 'j'); + std::vector initial_slices{{initial.data(), initial.size()}}; + ASSERT_TRUE(writer_->Put(key, initial_slices, DfsConfig()).has_value()); + ASSERT_TRUE(provider_->UnmountSegment(segment_, segment_size_).has_value()); + ExpectDfsOnlyValue(key, initial); + + WaitForReadLeaseExpiry(); + std::string updated(4096, 'k'); + std::vector updated_slices{{updated.data(), updated.size()}}; + auto mismatched = writer_->Upsert(key, updated_slices, DfsConfig()); + ASSERT_FALSE(mismatched.has_value()); + EXPECT_EQ(mismatched.error(), ErrorCode::INVALID_PARAMS); + ExpectDfsOnlyValue(key, initial); + + WaitForReadLeaseExpiry(); + ASSERT_TRUE( + writer_->Upsert(key, updated_slices, DfsOnlyConfig()).has_value()); + ExpectDfsOnlyValue(key, updated); +} + } // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index 04f2387804..79792d4ab5 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -578,6 +578,150 @@ TEST_F(MasterServiceTest, DfsPutEndAllAndUpsertTopologyAreAtomic) { std::filesystem::remove_all(dfs_root, ec); } +// A one-object shard allocator: any DFS extent leaked by a lifecycle path makes +// the next DFS-only allocation fail. +class OneSlotDfsShardEnv { + public: + explicit OneSlotDfsShardEnv(const std::string& root) + : enable_dfs_("MOONCAKE_ENABLE_DFS", "1"), + fs_adapter_("MOONCAKE_DFS_FS_ADAPTER", "posix"), + root_dir_("MOONCAKE_DFS_ROOT_DIR", root.c_str()), + shard_count_("MOONCAKE_DFS_SHARD_COUNT", "1"), + shard_capacity_("MOONCAKE_DFS_SHARD_CAPACITY", "32768"), + alignment_("MOONCAKE_DFS_ALIGNMENT", "4096"), + eviction_("MOONCAKE_DFS_EVICTION_ENABLED", "0"), + deferred_free_("MOONCAKE_DFS_DEFERRED_FREE_SECONDS", "0"), + single_tenant_("MOONCAKE_DFS_SINGLE_TENANT", "true") {} + + static constexpr uint64_t kObjectSize = 20480; + + private: + ScopedEnvVar enable_dfs_, fs_adapter_, root_dir_, shard_count_, + shard_capacity_, alignment_, eviction_, deferred_free_, single_tenant_; +}; + +ReplicateConfig DfsOnlyConfig() { + ReplicateConfig config; + config.replica_num = 0; + config.dfs_replica_num = 1; + return config; +} + +TEST_F(MasterServiceTest, DfsOnlyLifecyclePathsReleaseTheExtent) { + const auto dfs_root = (std::filesystem::temp_directory_path() / + ("master_dfs_only_" + std::to_string(::getpid()))) + .string(); + std::filesystem::create_directories(dfs_root); + OneSlotDfsShardEnv env(dfs_root); + constexpr uint64_t kSize = OneSlotDfsShardEnv::kObjectSize; + + MasterServiceConfig master_config; + master_config.put_start_discard_timeout_sec = 1; + master_config.put_start_release_timeout_sec = 2; + MasterService service(master_config); + const UUID writer = generate_uuid(); + const UUID other = generate_uuid(); + + auto start = service.PutStart(writer, "a", TenantId::Default(), kSize, + DfsOnlyConfig()); + ASSERT_TRUE(start.has_value()) << toString(start.error()); + ASSERT_EQ(start->size(), 1u); + EXPECT_TRUE(start->front().is_dfs_replica()); + auto not_ready = service.GetReplicaList("a", TenantId::Default()); + ASSERT_FALSE(not_ready.has_value()); + EXPECT_EQ(not_ready.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto full = service.PutStart(writer, "b", TenantId::Default(), kSize, + DfsOnlyConfig()); + ASSERT_FALSE(full.has_value()); + EXPECT_EQ(full.error(), ErrorCode::NO_AVAILABLE_HANDLE); + + // Revoke releases the extent. + ASSERT_TRUE( + service.PutRevoke(writer, "a", TenantId::Default(), ReplicaType::ALL)); + auto revoked = service.GetReplicaList("a", TenantId::Default()); + ASSERT_FALSE(revoked.has_value()); + EXPECT_EQ(revoked.error(), ErrorCode::OBJECT_NOT_FOUND); + + // Upsert preemption releases the in-flight writer's extent. + ASSERT_TRUE(service.PutStart(writer, "b", TenantId::Default(), kSize, + DfsOnlyConfig())); + ASSERT_TRUE(service.UpsertStart(other, "b", TenantId::Default(), kSize, + DfsOnlyConfig())); + auto stale_end = + service.PutEnd(writer, "b", TenantId::Default(), ReplicaType::ALL); + ASSERT_FALSE(stale_end.has_value()); + EXPECT_EQ(stale_end.error(), ErrorCode::ILLEGAL_CLIENT); + ASSERT_TRUE( + service.UpsertEnd(other, "b", TenantId::Default(), ReplicaType::ALL)); + auto complete = service.GetReplicaList("b", TenantId::Default()); + ASSERT_TRUE(complete.has_value()); + ASSERT_EQ(complete->replicas.size(), 1u); + EXPECT_EQ(complete->replicas[0].status, ReplicaStatus::COMPLETE); + + // Remove honours the read lease, then releases the extent. + auto leased = service.Remove("b", TenantId::Default()); + ASSERT_FALSE(leased.has_value()); + EXPECT_EQ(leased.error(), ErrorCode::OBJECT_HAS_LEASE); + ASSERT_TRUE(service.Remove("b", TenantId::Default(), /*force=*/true)); + + // An abandoned write is reaped after the release timeout. + ASSERT_TRUE(service.PutStart(writer, "c", TenantId::Default(), kSize, + DfsOnlyConfig())); + std::this_thread::sleep_for(std::chrono::seconds(4)); + auto reaped = service.GetReplicaList("c", TenantId::Default()); + ASSERT_FALSE(reaped.has_value()); + EXPECT_EQ(reaped.error(), ErrorCode::OBJECT_NOT_FOUND); + auto reuse = service.PutStart(other, "d", TenantId::Default(), kSize, + DfsOnlyConfig()); + ASSERT_TRUE(reuse.has_value()) << toString(reuse.error()); + + std::error_code ec; + std::filesystem::remove_all(dfs_root, ec); +} + +TEST_F(MasterServiceTest, DfsOnlyObjectsIgnoreMemoryEvictionAndSegmentLoss) { + const auto dfs_root = + (std::filesystem::temp_directory_path() / + ("master_dfs_only_eviction_" + std::to_string(::getpid()))) + .string(); + std::filesystem::create_directories(dfs_root); + OneSlotDfsShardEnv env(dfs_root); + + MasterServiceConfig master_config; + master_config.default_kv_lease_ttl = 1; + MasterService service(master_config); + const auto context = PrepareSimpleSegment(service); + const UUID writer = generate_uuid(); + + ASSERT_TRUE(service.PutStart(writer, "dfs_only", TenantId::Default(), + OneSlotDfsShardEnv::kObjectSize, + DfsOnlyConfig())); + ASSERT_TRUE(service.PutEnd(writer, "dfs_only", TenantId::Default(), + ReplicaType::ALL)); + ReplicateConfig memory_config; + memory_config.replica_num = 1; + for (int i = 0; i < 8; ++i) { + const auto key = "memory_" + std::to_string(i); + ASSERT_TRUE(service.PutStart(writer, key, TenantId::Default(), 1 << 20, + memory_config)); + ASSERT_TRUE(service.PutEnd(writer, key, TenantId::Default(), + ReplicaType::MEMORY)); + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + MasterServiceTestPeer(service).RunBatchEvictForTesting(1.0, 1.0); + + auto query = service.GetReplicaList("dfs_only", TenantId::Default()); + ASSERT_TRUE(query.has_value()) << toString(query.error()); + ASSERT_EQ(query->replicas.size(), 1u); + EXPECT_TRUE(query->replicas[0].is_dfs_replica()); + ASSERT_TRUE(service.UnmountSegment(context.segment_id, context.client_id)); + EXPECT_TRUE( + service.GetReplicaList("dfs_only", TenantId::Default()).has_value()); + + std::error_code ec; + std::filesystem::remove_all(dfs_root, ec); +} + TEST_F(MasterServiceTest, DfsBucketMemoryAllocationFailurePreservesBuckets) { const auto dfs_root = (std::filesystem::temp_directory_path() / From 3cf058db6d4237b1fc3e63d1068430d3b0983bc8 Mon Sep 17 00:00:00 2001 From: Zicai Pu <102817879+pzcyyj@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:02:26 +0800 Subject: [PATCH 10/14] [Store] Add opt-in same-cycle flush for partial SSD offload buckets (#4673) --- docs/source/design/store/ssd-offload.md | 4 + .../include/environment_variables.h | 1 + .../include/config/bucket_backend_config.h | 4 + .../src/config/bucket_backend_config.cpp | 4 + mooncake-store/src/storage_backend.cpp | 8 ++ .../tests/bucket_backend_config_test.cpp | 29 +++++ mooncake-store/tests/storage_backend_test.cpp | 106 ++++++++++++++++++ 7 files changed, 156 insertions(+) diff --git a/docs/source/design/store/ssd-offload.md b/docs/source/design/store/ssd-offload.md index adc952b8f4..06c4b7bc43 100644 --- a/docs/source/design/store/ssd-offload.md +++ b/docs/source/design/store/ssd-offload.md @@ -144,6 +144,10 @@ Bucket IDs are monotonically increasing timestamps with a sequence suffix, so `b **Grouping strategy** (`GroupOffloadingKeysByBucket`): objects are accumulated into a bucket until either `bucket_size_limit` (default 256 MB) or `bucket_keys_limit` (default 500) is reached. Objects that do not fill a complete bucket are held in `ungrouped_offloading_objects_` and retried on the next heartbeat. +Set `MOONCAKE_OFFLOAD_FLUSH_PARTIAL_BUCKETS=1` before starting a client to write its final partial bucket in the same offload cycle. The default is `0`, preserving the existing grouping policy. The setting is read when the bucket backend configuration is created; an existing backend does not change when the environment changes. + +This opt-in policy avoids the deferred-task loss described in [issue #3006](https://github.com/kvcache-ai/Mooncake/issues/3006): carrying only keys and sizes across heartbeats does not preserve their offload tasks. Flushing the tail keeps eligible objects in the cycle that owns those tasks. It does not change completion notification, write-failure handling, or the on-disk format. Smaller buckets can increase file/metadata overhead and reduce write efficiency; this policy does not guarantee a disk replica if an offload fails. An alternative is to retain the tasks across cycles while waiting for full buckets, as proposed in [PR #3009](https://github.com/kvcache-ai/Mooncake/pull/3009). + **In-flight read tracking**: A `BucketReadGuard` RAII object increments `BucketMetadata::inflight_reads_` on construction and decrements it on destruction. This allows safe deletion of bucket files even when concurrent reads are in progress. **Ownership contract (`storage_path`)**: LOCAL_DISK allows **one live client per `storage_path`**. `BucketStorageBackend::Init()` opens `.mooncake_local_disk.lock` under that directory and takes a non-blocking exclusive `flock`; a second live client fails Init immediately with an actionable error (`storage_path already held by another live client … use standalone store or per-client directories`). Do not delete the lock file while a client is running β€” that would release the flock for other processes. `flock` is only meaningful on local filesystems (NFS and similar may not honor exclusive locks the same way). diff --git a/mooncake-common/include/environment_variables.h b/mooncake-common/include/environment_variables.h index 1d48347201..f52c10e616 100644 --- a/mooncake-common/include/environment_variables.h +++ b/mooncake-common/include/environment_variables.h @@ -45,6 +45,7 @@ struct FilePerKeyEnvironmentVariables { }; struct BucketBackendEnvironmentVariables { + MC_DEFINE_ENV_VAR(bool, MOONCAKE_OFFLOAD_FLUSH_PARTIAL_BUCKETS); MC_DEFINE_ENV_VAR(int64_t, MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT); MC_DEFINE_ENV_VAR(int64_t, MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES); MC_DEFINE_ENV_VAR(int64_t, MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE); diff --git a/mooncake-store/include/config/bucket_backend_config.h b/mooncake-store/include/config/bucket_backend_config.h index 17f04d7d5b..1007d59b62 100644 --- a/mooncake-store/include/config/bucket_backend_config.h +++ b/mooncake-store/include/config/bucket_backend_config.h @@ -62,6 +62,10 @@ struct BucketBackendConfig { // max_physical_bytes > 0. int64_t disk_scan_cache_ms = 500; + // Write the final partial bucket in the current offload cycle instead of + // carrying keys into a later cycle without their current offload tasks. + bool flush_partial_buckets = false; + bool Validate() const; static BucketBackendConfig FromEnvironment(); diff --git a/mooncake-store/src/config/bucket_backend_config.cpp b/mooncake-store/src/config/bucket_backend_config.cpp index 4417d596fc..f6c84a8c29 100644 --- a/mooncake-store/src/config/bucket_backend_config.cpp +++ b/mooncake-store/src/config/bucket_backend_config.cpp @@ -25,6 +25,10 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { BucketBackendConfig config; using Variables = BucketBackendEnvironmentVariables; + config.flush_partial_buckets = + Environ::ReadOr(Variables::MOONCAKE_OFFLOAD_FLUSH_PARTIAL_BUCKETS, + config.flush_partial_buckets); + config.bucket_keys_limit = Environ::ReadOr(Variables::MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT, config.bucket_keys_limit); diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index bde6e254e9..130f5f0324 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -2347,6 +2347,14 @@ tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( while (static_cast(bucket_keys.size()) < bucket_backend_config_.bucket_keys_limit) { if (it == offloading_objects.cend()) { + if (bucket_backend_config_.flush_partial_buckets) { + // Keep these keys in the current cycle, while FileStorage + // still has their tasks and source replicas are pinned. + if (!bucket_keys.empty()) { + buckets_keys.push_back(std::move(bucket_keys)); + } + return {}; + } for (const auto& bucket_object : bucket_objects) { ungrouped_offloading_objects.emplace(bucket_object.first, bucket_object.second); diff --git a/mooncake-store/tests/bucket_backend_config_test.cpp b/mooncake-store/tests/bucket_backend_config_test.cpp index 3991508657..04a6bec8b3 100644 --- a/mooncake-store/tests/bucket_backend_config_test.cpp +++ b/mooncake-store/tests/bucket_backend_config_test.cpp @@ -49,6 +49,8 @@ struct BucketBackendEnvironment { "MOONCAKE_OFFLOAD_BUCKET_DISK_SCAN_CACHE_MS"}; ScopedEnvVar eviction_policy{"MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY"}; ScopedEnvVar legacy_eviction_policy{"MOONCAKE_BUCKET_EVICTION_POLICY"}; + ScopedEnvVar flush_partial_buckets{ + "MOONCAKE_OFFLOAD_FLUSH_PARTIAL_BUCKETS"}; }; class BucketBackendConfigTest : public ::testing::Test { @@ -65,6 +67,33 @@ TEST_F(BucketBackendConfigTest, UsesExistingDefaultsWhenEnvironmentIsUnset) { EXPECT_EQ(config.max_physical_bytes, 0); EXPECT_EQ(config.disk_scan_cache_ms, 500); EXPECT_EQ(config.eviction_policy, BucketEvictionPolicy::FIFO); + EXPECT_FALSE(config.flush_partial_buckets); +} + +TEST_F(BucketBackendConfigTest, ReadsPartialFlushBooleanValues) { + struct Case { + const char* value; + bool enabled; + }; + for (const auto& entry : + {Case{"1", true}, Case{"true", true}, Case{"0", false}, + Case{"FALSE", false}, Case{"invalid", false}}) { + SCOPED_TRACE(entry.value); + env.flush_partial_buckets.Set(entry.value); + EXPECT_EQ(BucketBackendConfig::FromEnvironment().flush_partial_buckets, + entry.enabled); + } +} + +TEST_F(BucketBackendConfigTest, PartialFlushIsReadPerConfiguration) { + env.flush_partial_buckets.Set("1"); + const auto enabled = BucketBackendConfig::FromEnvironment(); + EXPECT_TRUE(enabled.flush_partial_buckets); + env.flush_partial_buckets.Set("0"); + EXPECT_FALSE(BucketBackendConfig::FromEnvironment().flush_partial_buckets); + env.flush_partial_buckets.Unset(); + EXPECT_FALSE(BucketBackendConfig::FromEnvironment().flush_partial_buckets); + EXPECT_TRUE(enabled.flush_partial_buckets); } TEST_F(BucketBackendConfigTest, ReadsIndependentValidValues) { diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 8c69db6448..eb9c7fe874 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -2138,6 +2138,112 @@ TEST_F(StorageBackendTest, // BucketStorageBackend: Duplicate Key Detection Tests (Phase 0 - D0) //----------------------------------------------------------------------------- +TEST_F(StorageBackendTest, + BucketStorageBackend_PartialFlushDisabledRetainsTail) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 4; + bucket_config.bucket_size_limit = 64; + EXPECT_FALSE(bucket_config.flush_partial_buckets); + BucketStorageBackend backend(config, bucket_config); + ASSERT_TRUE(backend.Init()); + std::vector> buckets; + ASSERT_TRUE( + backend.AllocateOffloadingBuckets({{"A1", 16}, {"A2", 16}}, buckets)); + EXPECT_TRUE(buckets.empty()); + EXPECT_EQ(backend.UngroupedOffloadingObjectsSize(), 2); + ASSERT_TRUE( + backend.AllocateOffloadingBuckets({{"B1", 16}, {"B2", 16}}, buckets)); + ASSERT_EQ(buckets.size(), 1); + std::sort(buckets[0].begin(), buckets[0].end()); + EXPECT_EQ(buckets[0], (std::vector{"A1", "A2", "B1", "B2"})); + EXPECT_EQ(backend.UngroupedOffloadingObjectsSize(), 0); +} + +TEST_F(StorageBackendTest, + BucketStorageBackend_PartialFlushKeepsCyclesSeparate) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 4; + bucket_config.bucket_size_limit = 64; + bucket_config.flush_partial_buckets = true; + BucketStorageBackend backend(config, bucket_config); + ASSERT_TRUE(backend.Init()); + for (const auto* prefix : {"A", "B"}) { + const std::string first = std::string(prefix) + "1"; + const std::string second = std::string(prefix) + "2"; + std::vector> buckets; + ASSERT_TRUE(backend.AllocateOffloadingBuckets( + {{first, 16}, {second, 16}}, buckets)); + ASSERT_EQ(buckets.size(), 1); + std::sort(buckets[0].begin(), buckets[0].end()); + EXPECT_EQ(buckets[0], (std::vector{first, second})); + EXPECT_EQ(backend.UngroupedOffloadingObjectsSize(), 0); + } +} + +TEST_F(StorageBackendTest, BucketStorageBackend_PartialFlushSkipsEmptyBuckets) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 4; + bucket_config.bucket_size_limit = 64; + bucket_config.flush_partial_buckets = true; + BucketStorageBackend backend(config, bucket_config); + ASSERT_TRUE(backend.Init()); + std::string value = "value"; + std::unordered_map> batch; + batch.emplace("existing", std::vector{{value.data(), value.size()}}); + ASSERT_TRUE(backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; })); + using Objects = std::unordered_map; + for (const auto& objects : + {Objects{}, Objects{{"existing", 5}}, Objects{{"oversized", 65}}, + Objects{{"existing", 5}, {"oversized", 65}}}) { + std::vector> buckets; + ASSERT_TRUE(backend.AllocateOffloadingBuckets(objects, buckets)); + EXPECT_TRUE(buckets.empty()); + EXPECT_EQ(backend.UngroupedOffloadingObjectsSize(), 0); + } +} + +TEST_F(StorageBackendTest, BucketStorageBackend_PartialFlushRespectsLimits) { + for (const auto& limits : {std::pair{2, 64}, + std::pair{8, 24}}) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = limits.first; + bucket_config.bucket_size_limit = limits.second; + bucket_config.flush_partial_buckets = true; + BucketStorageBackend backend(config, bucket_config); + ASSERT_TRUE(backend.Init()); + const std::unordered_map objects{ + {"A", 8}, {"B", 8}, {"C", 8}, {"D", 8}, {"E", 8}}; + std::vector> buckets; + ASSERT_TRUE(backend.AllocateOffloadingBuckets(objects, buckets)); + EXPECT_EQ(buckets.size(), limits.first == 2 ? 3 : 2); + std::vector all_keys; + for (const auto& bucket : buckets) { + EXPECT_FALSE(bucket.empty()); + EXPECT_LE(bucket.size(), + static_cast(bucket_config.bucket_keys_limit)); + int64_t bytes = 0; + for (const auto& key : bucket) bytes += objects.at(key); + EXPECT_LE(bytes, bucket_config.bucket_size_limit); + all_keys.insert(all_keys.end(), bucket.begin(), bucket.end()); + } + std::sort(all_keys.begin(), all_keys.end()); + EXPECT_EQ(all_keys, + (std::vector{"A", "B", "C", "D", "E"})); + EXPECT_EQ(backend.UngroupedOffloadingObjectsSize(), 0); + } +} + TEST_F(StorageBackendTest, BucketStorageBackend_DuplicateKeyIdempotentSkip) { FileStorageConfig config; config.storage_filepath = data_path; From 292e8391b7df0a385092474374660f599958f6c4 Mon Sep 17 00:00:00 2001 From: tensorcopy Date: Mon, 19 Oct 2026 00:21:32 -0700 Subject: [PATCH 11/14] [Store] HA: make the etcd leader lease TTL configurable (#4706) The leadership lease TTL was hard-coded to DEFAULT_MASTER_VIEW_LEASE_TTL_SEC (5 s) in the etcd coordinator. An etcd apply stall of about that length makes the leader miss a keepalive and step down, so clusters with observed slow applies cannot ride such stalls out. Add --ha_leader_lease_ttl_sec (default 5, allowed range 2..60, also read from the JSON config as ha_leader_lease_ttl_sec) and thread it through MasterConfig, MasterServiceSupervisorConfig and HABackendSpec to the etcd coordinator's GrantLease call. The effective value is logged at startup. The floor is 2 s because etcd raises smaller lease TTLs to 2 s with its default settings. The default keeps the existing behavior; the Redis and k8s backends do not use the flag. Document the flag in the Mooncake Store deployment guide. A longer TTL trades failover latency for stall tolerance: a new leader still waits one full TTL in warmup before serving (fencing against the old leader), and after a leader crash the standby waits for the old lease to expire. --- .../mooncake-store-deployment-guide.md | 3 +- mooncake-store/include/ha/ha_types.h | 2 + mooncake-store/include/master_config.h | 3 + .../backends/etcd/etcd_leader_coordinator.cpp | 7 +- .../leadership/master_service_supervisor.cpp | 1 + mooncake-store/src/master.cpp | 32 +++ .../leadership/etcd_leader_lease_ttl_e2e.sh | 193 ++++++++++++++++++ mooncake-store/tests/master_config_test.cpp | 19 ++ 8 files changed, 257 insertions(+), 3 deletions(-) create mode 100755 mooncake-store/tests/ha/leadership/etcd_leader_lease_ttl_e2e.sh diff --git a/docs/source/deployment/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index d185ce7790..d3fde8da15 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -361,7 +361,7 @@ With `ha_backend_type=etcd`, Mooncake Store supports a Primary-Standby HA model HA leadership and metadata replication are configured separately: - The HA coordinator elects the active master. Configure it with `--enable_ha`, `--ha_backend_type`, `--ha_backend_connstring`, and `--cluster_id`. For `ha_backend_type=etcd`, legacy `--etcd_endpoints` is used only when `--ha_backend_connstring` is empty. -- The optional batch-record OpLog records metadata mutations so standby masters can catch up and later be promoted. It is the etcd Primary-Standby path. Enable it explicitly with `--enable_oplog=true`; it is disabled by default. `ha_backend_type=etcd` requires a build with `STORE_USE_ETCD`. On etcd, acknowledged batches form the durable prefix used for promotion. `ha_backend_type=redis` requires a build with `STORE_USE_REDIS` and is an experimental HA log storage method. A successful Redis write only means Redis accepted the record. Do not treat Redis persistence, including AOF `appendfsync everysec`, as a strict durability guarantee. +- The optional batch-record OpLog records metadata mutations so standby masters can catch up and later be promoted. It is the etcd Primary-Standby path. Enable it explicitly with `--enable_oplog=true`; it is disabled by default. `ha_backend_type=etcd` requires a build with `STORE_USE_ETCD`. On etcd, acknowledged batches form the durable prefix used for promotion. `ha_backend_type=redis` requires a build with `STORE_USE_REDIS` and is an experimental HA log storage method. A successful Redis write only means Redis accepted the record. Do not treat Redis persistence, including AOF `appendfsync everysec`, as a strict durability guarantee. - The optional standby-generated batch OpLog snapshot path is enabled with `--enable_oplog_snapshot=true` together with `--enable_oplog=true`. It uses the batch snapshot provider/coordinator and does not use the legacy catalog snapshot manager. Startup fails when the required etcd, cluster ID, object-store, or chunk configuration is invalid; a temporary upload failure leaves OpLog apply running for a later attempt. @@ -745,6 +745,7 @@ mooncake_master \ | `--ha_backend_connstring` | empty | HA backend connection string | | `--etcd_endpoints` | empty | Backward-compatible etcd HA endpoints, used only for `ha_backend_type=etcd` when `--ha_backend_connstring` is empty | | `--cluster_id` | `mooncake_cluster` | Cluster ID for HA persistence | +| `--ha_leader_lease_ttl_sec` | `5` | Leadership lease TTL in seconds for `ha_backend_type=etcd` (range `2`-`60`; the JSON config key `ha_leader_lease_ttl_sec` is also accepted). Raise it to ride out etcd apply stalls longer than the default. A longer TTL delays failover: a new leader waits one full TTL before serving, and a standby waits for a crashed leader's lease to expire. Ignored by the Redis and k8s backends | | `--enable_oplog` | `false` | Enable the batch OpLog writer and reader; requires `enable_ha=true` and `ha_backend_type=etcd` or `redis`. etcd uses it for Primary-Standby promotion. Redis uses it as an experimental HA log store: a successful write only means Redis accepted the record | | `--enable_oplog_snapshot` | `false` | Enable standby-generated batch OpLog snapshots; requires batch OpLog, HA/etcd, valid object-store configuration, and persistent local snapshot storage when applicable | | `--snapshot_chunk_object_count` | `1000000` | Maximum objects per batch OpLog snapshot chunk; must be positive when the new snapshot path is enabled | diff --git a/mooncake-store/include/ha/ha_types.h b/mooncake-store/include/ha/ha_types.h index d4506fb646..0b6ca7432a 100644 --- a/mooncake-store/include/ha/ha_types.h +++ b/mooncake-store/include/ha/ha_types.h @@ -99,6 +99,8 @@ struct HABackendSpec { HABackendType type = HABackendType::UNKNOWN; std::string connstring; ClusterNamespace cluster_namespace; + // Leadership lease TTL in seconds. Currently honored by the etcd backend. + int64_t leader_lease_ttl_sec = DEFAULT_MASTER_VIEW_LEASE_TTL_SEC; }; struct MasterView { diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 6c22d8573b..489d02a616 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -121,6 +121,7 @@ struct MasterConfig { bool enable_offload; std::string ha_backend_type; std::string ha_backend_connstring; + int64_t ha_leader_lease_ttl_sec = DEFAULT_MASTER_VIEW_LEASE_TTL_SEC; std::string etcd_endpoints; // OpLog store configuration @@ -257,6 +258,7 @@ class MasterServiceSupervisorConfig { uint32_t handler_worker_thread_num = 0; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; + int64_t ha_leader_lease_ttl_sec = DEFAULT_MASTER_VIEW_LEASE_TTL_SEC; std::string etcd_endpoints = "0.0.0.0:2379"; // OpLog store configuration bool enable_oplog = false; @@ -393,6 +395,7 @@ class MasterServiceSupervisorConfig { rpc_enable_tcp_no_delay = config.rpc_enable_tcp_no_delay; handler_worker_thread_num = config.handler_worker_thread_num; ha_backend_type = config.ha_backend_type; + ha_leader_lease_ttl_sec = config.ha_leader_lease_ttl_sec; etcd_endpoints = config.etcd_endpoints; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, etcd_endpoints); diff --git a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp index 40a875baf7..2e4e7073bd 100644 --- a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp +++ b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp @@ -213,8 +213,11 @@ EtcdLeaderCoordinator::TryAcquireLeadership(const std::string& leader_address) { return tl::make_unexpected(err); } + const int64_t lease_ttl_sec = spec_.leader_lease_ttl_sec > 0 + ? spec_.leader_lease_ttl_sec + : DEFAULT_MASTER_VIEW_LEASE_TTL_SEC; EtcdLeaseId lease_id = 0; - err = EtcdHelper::GrantLease(DEFAULT_MASTER_VIEW_LEASE_TTL_SEC, lease_id); + err = EtcdHelper::GrantLease(lease_ttl_sec, lease_id); if (err != ErrorCode::OK) { (void)ResetConnection(); return tl::make_unexpected(err); @@ -257,7 +260,7 @@ EtcdLeaderCoordinator::TryAcquireLeadership(const std::string& leader_address) { .view = MasterView{.leader_address = leader_address, .view_version = view_version}, .owner_token = owner_token, - .lease_ttl = std::chrono::seconds(DEFAULT_MASTER_VIEW_LEASE_TTL_SEC), + .lease_ttl = std::chrono::seconds(lease_ttl_sec), }; { diff --git a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp index 7a4e1a8c0c..a6c766bfe7 100644 --- a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp +++ b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp @@ -82,6 +82,7 @@ tl::expected BuildHABackendSpec( .type = backend_type.value(), .connstring = connstring, .cluster_namespace = config.cluster_id, + .leader_lease_ttl_sec = config.ha_leader_lease_ttl_sec, }; } diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 8f2446fc8b..65aa1adf42 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -51,6 +51,9 @@ static_assert(mooncake::DEFAULT_MAX_KV_SOFT_PIN_TTL_MS == 24 * 60 * 60 * 1000, constexpr char kDefaultKvLeaseTtlFlagValue[] = "10000"; constexpr char kDefaultKvSoftPinTtlFlagValue[] = "1800000"; constexpr char kDefaultMaxKvSoftPinTtlFlagValue[] = "86400000"; +// Allowed range of --ha_leader_lease_ttl_sec. +constexpr int64_t kMinHaLeaderLeaseTtlSec = 2; +constexpr int64_t kMaxHaLeaderLeaseTtlSec = 60; namespace { @@ -322,6 +325,20 @@ DEFINE_uint32(kv_events_queue_capacity, 65536, "the queue is full (0 = unbounded)"); DEFINE_string(ha_backend_type, "etcd", "HA backend type, e.g. etcd | redis | k8s"); +DEFINE_int64(ha_leader_lease_ttl_sec, + mooncake::DEFAULT_MASTER_VIEW_LEASE_TTL_SEC, + "Leadership lease TTL in seconds for the etcd HA backend. A " + "longer TTL tolerates longer etcd apply stalls but delays " + "failover after a leader crash (range 2..60)"); +DEFINE_validator(ha_leader_lease_ttl_sec, [](const char* flagname, + int64_t value) { + if (value < kMinHaLeaderLeaseTtlSec || value > kMaxHaLeaderLeaseTtlSec) { + LOG(ERROR) << flagname << " must be in [" << kMinHaLeaderLeaseTtlSec + << ", " << kMaxHaLeaderLeaseTtlSec << "], got " << value; + return false; + } + return true; +}); DEFINE_string(ha_backend_connstring, "", "HA backend connection string. If unset, only backend_type=etcd " "falls back to etcd_endpoints for backward compatibility"); @@ -655,6 +672,9 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetString("ha_backend_connstring", &master_config.ha_backend_connstring, FLAGS_ha_backend_connstring); + default_config.GetInt64("ha_leader_lease_ttl_sec", + &master_config.ha_leader_lease_ttl_sec, + FLAGS_ha_leader_lease_ttl_sec); default_config.GetString("etcd_endpoints", &master_config.etcd_endpoints, FLAGS_etcd_endpoints); default_config.GetString("cluster_id", &master_config.cluster_id, @@ -1085,6 +1105,17 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, !conf_set) { master_config.ha_backend_connstring = FLAGS_ha_backend_connstring; } + if ((google::GetCommandLineFlagInfo("ha_leader_lease_ttl_sec", &info) && + !info.is_default) || + !conf_set) { + master_config.ha_leader_lease_ttl_sec = FLAGS_ha_leader_lease_ttl_sec; + } + if (master_config.ha_leader_lease_ttl_sec < kMinHaLeaderLeaseTtlSec || + master_config.ha_leader_lease_ttl_sec > kMaxHaLeaderLeaseTtlSec) { + LOG(FATAL) << "ha_leader_lease_ttl_sec must be in [" + << kMinHaLeaderLeaseTtlSec << ", " << kMaxHaLeaderLeaseTtlSec + << "], got " << master_config.ha_leader_lease_ttl_sec; + } if ((google::GetCommandLineFlagInfo("etcd_endpoints", &info) && !info.is_default) || !conf_set) { @@ -1569,6 +1600,7 @@ int main(int argc, char* argv[]) { << ", offload_cap_ratio=" << master_config.offload_cap_ratio << ", ha_backend_type=" << master_config.ha_backend_type << ", ha_backend_connstring=" << ha_backend_connstring + << ", ha_leader_lease_ttl_sec=" << master_config.ha_leader_lease_ttl_sec << ", etcd_endpoints=" << master_config.etcd_endpoints << ", client_active_ttl_sec=" << master_config.client_active_ttl_sec << ", client_suspicion_ttl_sec=" diff --git a/mooncake-store/tests/ha/leadership/etcd_leader_lease_ttl_e2e.sh b/mooncake-store/tests/ha/leadership/etcd_leader_lease_ttl_e2e.sh new file mode 100755 index 0000000000..2ad5d920bb --- /dev/null +++ b/mooncake-store/tests/ha/leadership/etcd_leader_lease_ttl_e2e.sh @@ -0,0 +1,193 @@ +#!/usr/bin/env bash +# Checks that --ha_leader_lease_ttl_sec and the JSON key ha_leader_lease_ttl_sec +# reach the lease a running HA master is granted, against a local one-member +# etcd (same approach as etcd_leader_hang_e2e.sh): +# - no flag: the default TTL (5 s) is granted; +# - 15 and the range boundaries 2 and 60 are granted exactly; +# - the JSON config key overrides the default; +# - out-of-range values (1 and 61, flag or JSON) make the master exit +# nonzero with the validation error, before it connects to etcd. +# A master that exits during a success case fails the run. +# +# Requires mooncake_master from a normal build (no failpoints needed). +set -euo pipefail +# Case 8 refuses a bad JSON value through LOG(FATAL); do not leave a core file. +ulimit -c 0 + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +MASTER_BIN="${MASTER_BIN:-${ROOT_DIR}/build/mooncake-store/src/mooncake_master}" +ETCD_BIN="${ETCD_BIN:-etcd}" +ETCDCTL_BIN="${ETCDCTL_BIN:-etcdctl}" +# All ports stay inside the local test range 23000-23999; every case uses a +# base of 231xx: etcd client base+1, peer base+2, master rpc base+20. +BASE=23100 + +work="$(mktemp -d)" +etcd_pid="" +master_pid="" +current_log="" + +cleanup() { + rc=$? + set +e + kill "${master_pid:-}" 2>/dev/null + wait "${master_pid:-}" 2>/dev/null + kill "${etcd_pid:-}" 2>/dev/null + wait "${etcd_pid:-}" 2>/dev/null + if [[ "${KEEP_WORK:-0}" != 1 ]]; then rm -rf "$work"; fi + exit "$rc" +} +trap cleanup EXIT + +fail() { + echo "FAIL: $*" >&2 + [[ -n "$current_log" ]] && tail -n 30 "$current_log" 2>/dev/null >&2 + exit 1 +} + +start_etcd() { + local port=$1 peer=$2 data=$3 + "$ETCD_BIN" --name ttl-check-etcd --data-dir "$data" \ + --listen-client-urls "http://127.0.0.1:${port}" \ + --advertise-client-urls "http://127.0.0.1:${port}" \ + --listen-peer-urls "http://127.0.0.1:${peer}" \ + --initial-advertise-peer-urls "http://127.0.0.1:${peer}" \ + --initial-cluster ttl-check-etcd=http://127.0.0.1:${peer} \ + >"$work/etcd.log" 2>&1 & + etcd_pid=$! + for _ in $(seq 1 30); do + "$ETCDCTL_BIN" --endpoints="127.0.0.1:${port}" --command-timeout=1s \ + endpoint health >/dev/null 2>&1 && return 0 + sleep 1 + done + fail "etcd on ${port} did not become healthy" +} + +start_master() { + local rpc=$1 log=$2 + shift 2 + current_log="$log" + "$MASTER_BIN" "$@" \ + --rpc_address=127.0.0.1 --rpc_port="$rpc" \ + --logtostderr=1 >"$log" 2>&1 & + master_pid=$! +} + +view_key() { echo "mooncake-store/${CLUSTER_ID}/master_view"; } + +# Reads the master_view lease id and granted TTL from etcd. Only valid +# while the master is running. +read_granted_ttl() { + local lease granted + lease=$("$ETCDCTL_BIN" --endpoints="127.0.0.1:${ETCD_PORT}" get \ + "$(view_key)" -w json 2>/dev/null | python3 -c \ + 'import json,sys; kvs=json.load(sys.stdin).get("kvs",[]); print(kvs[0].get("lease",0) if kvs else 0)') + [[ "$lease" != 0 ]] || fail "master_view has no lease attached" + granted=$("$ETCDCTL_BIN" --endpoints="127.0.0.1:${ETCD_PORT}" lease \ + timetolive "$(printf '%x' "$lease")" -w json 2>/dev/null | python3 -c \ + 'import json,sys; print(json.load(sys.stdin)["granted-ttl"])') + echo "$granted" +} + +stop_master() { + kill "${master_pid:-}" 2>/dev/null || true + wait "${master_pid:-}" 2>/dev/null || true + master_pid="" +} + +stop_etcd() { + kill "${etcd_pid:-}" 2>/dev/null || true + wait "${etcd_pid:-}" 2>/dev/null || true + etcd_pid="" + rm -rf "$work/etcd-data" +} + +case_no=0 +next_case() { + case_no=$((case_no + 1)) + ETCD_PORT=$((BASE + case_no * 10 + 1)) + stop_etcd + start_etcd "$ETCD_PORT" $((ETCD_PORT + 1)) "$work/etcd-data" + CLUSTER_ID="ttl-e2e-case-${case_no}-$$" + mkdir -p "$work/case$case_no" +} + +# Expected-success case: the master publishes its view under a lease whose +# granted TTL must be exactly $1. Any master exit during the wait fails. +run_granted_case() { + local expect_ttl=$1 + shift + next_case + local rpc=$((BASE + case_no * 10 + 20)) log="$work/case$case_no/master.log" + start_master "$rpc" "$log" \ + --enable_ha=true --ha_backend_type=etcd \ + --etcd_endpoints="127.0.0.1:${ETCD_PORT}" --cluster_id="${CLUSTER_ID}" "$@" + local granted="" + for _ in $(seq 1 90); do + kill -0 "$master_pid" 2>/dev/null || + fail "case $case_no: master exited before publishing its view" + if "$ETCDCTL_BIN" --endpoints="127.0.0.1:${ETCD_PORT}" get "$(view_key)" \ + --print-value-only 2>/dev/null | grep -q "127.0.0.1:${rpc}"; then + granted=$(read_granted_ttl) + break + fi + sleep 1 + done + [[ -n "$granted" ]] || fail "case $case_no: no view published in 90 s" + [[ "$granted" == "$expect_ttl" ]] || + fail "case $case_no: granted TTL $granted, expected $expect_ttl" + grep -q "ha_leader_lease_ttl_sec=$expect_ttl" "$log" || + fail "case $case_no does not log the effective TTL" + echo "case $case_no: granted TTL $granted - PASS" + stop_master +} + +# Expected-rejection case: the out-of-range value is refused by config +# validation before the supervisor connects to etcd (gflags validator at +# ParseCommandLineFlags time for the flag; the range check in +# LoadConfigFromCmdline for the JSON key). The contract asserted is a +# nonzero exit plus the validation error in the log. +run_refused_case() { + next_case + local rpc=$((BASE + case_no * 10 + 20)) log="$work/case$case_no/master.log" + start_master "$rpc" "$log" \ + --enable_ha=true --ha_backend_type=etcd \ + --etcd_endpoints="127.0.0.1:${ETCD_PORT}" --cluster_id="${CLUSTER_ID}" "$@" + local rc="" + # Bounded wait for the refusal to complete: a validation exit is + # immediate, so the bound only fails a wedged binary. + for _ in $(seq 1 30); do + kill -0 "$master_pid" 2>/dev/null || break + sleep 1 + done + kill -0 "$master_pid" 2>/dev/null && + fail "case $case_no: out-of-range TTL was accepted (master still running)" + wait "$master_pid" 2>/dev/null || rc=$? + master_pid="" + [[ -n "$rc" && "$rc" != 0 ]] || + fail "case $case_no: out-of-range TTL exited ${rc:-0} instead of refusing" + grep -q "ha_leader_lease_ttl_sec must be in \[2, 60\]" "$log" || + fail "case $case_no: refused master did not log the validation error" + echo "case $case_no: refused (exit $rc) with the validation error logged - PASS" +} + +# 1. Default TTL: no flag -> granted 5. +run_granted_case 5 + +# 2. Explicit valid values incl. both range boundaries. +for ttl in 15 2 60; do + run_granted_case "$ttl" --ha_leader_lease_ttl_sec="$ttl" +done + +# 3. JSON config value (config_path) reaches the granted lease. +printf '{"ha_leader_lease_ttl_sec": 17}\n' >"$work/config.json" +run_granted_case 17 --config_path="$work/config.json" + +# 4. Out-of-range values (flag and JSON): refused at config validation. +for bad in 1 61; do + run_refused_case --ha_leader_lease_ttl_sec="$bad" +done +printf '{"ha_leader_lease_ttl_sec": 61}\n' >"$work/config-bad.json" +run_refused_case --config_path="$work/config-bad.json" + +echo "lease TTL propagation check passed in ${SECONDS}s" diff --git a/mooncake-store/tests/master_config_test.cpp b/mooncake-store/tests/master_config_test.cpp index 2e676ad763..52c7e134f5 100644 --- a/mooncake-store/tests/master_config_test.cpp +++ b/mooncake-store/tests/master_config_test.cpp @@ -61,5 +61,24 @@ TEST(ClientLivenessConfigTest, LegacyBuilderSetterDefaultsBothWindows) { EXPECT_EQ(config.client_suspicion_ttl_sec, 45); } +TEST(HALeaseTtlConfigTest, SupervisorConfigDefaultsToMasterViewTtl) { + const MasterServiceSupervisorConfig config{}; + + EXPECT_EQ(DEFAULT_MASTER_VIEW_LEASE_TTL_SEC, + config.ha_leader_lease_ttl_sec); +} + +TEST(HALeaseTtlConfigTest, SupervisorConfigCarriesMasterConfigTtl) { + MasterConfig master_config{}; + master_config.ha_backend_type = "etcd"; + EXPECT_EQ(DEFAULT_MASTER_VIEW_LEASE_TTL_SEC, + master_config.ha_leader_lease_ttl_sec); + master_config.ha_leader_lease_ttl_sec = 17; + + const MasterServiceSupervisorConfig supervisor_config(master_config); + + EXPECT_EQ(17, supervisor_config.ha_leader_lease_ttl_sec); +} + } // namespace } // namespace mooncake::test From 1a6d00d0f4dc7f54c573c5fcd4ff5f80bd1f2813 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Mon, 19 Oct 2026 15:43:51 +0800 Subject: [PATCH 12/14] [Bugfix] Snapshot file birth times before reconstructing eviction order (#4688) Signed-off-by: Xucheng Zhou --- mooncake-store/include/common/file_util.h | 5 + mooncake-store/src/common/file_util.cpp | 49 +++++++++ mooncake-store/src/storage_backend.cpp | 18 +--- mooncake-store/tests/file_util_test.cpp | 126 ++++++++++++++++++++++ 4 files changed, 181 insertions(+), 17 deletions(-) diff --git a/mooncake-store/include/common/file_util.h b/mooncake-store/include/common/file_util.h index 67388d640a..e94926a2f4 100644 --- a/mooncake-store/include/common/file_util.h +++ b/mooncake-store/include/common/file_util.h @@ -12,6 +12,11 @@ namespace fs = std::filesystem; class FileUtil { public: + // Snapshot birth times before sorting, oldest known first. Failed metadata + // queries and filesystems without birth times sort last, by path. Entries + // are retained so callers can separately check whether they are readable. + static void SortByCreationTime(std::vector& files); + static std::string ResolvePathFromKey(const std::string& key, const std::string& root_dir, const std::string& fsdir); diff --git a/mooncake-store/src/common/file_util.cpp b/mooncake-store/src/common/file_util.cpp index be10319982..cc8840d2de 100644 --- a/mooncake-store/src/common/file_util.cpp +++ b/mooncake-store/src/common/file_util.cpp @@ -1,10 +1,16 @@ #include "common/file_util.h" #include +#include #include #include #include #include #include +#include +#include + +#include +#include namespace mooncake { namespace { @@ -65,6 +71,49 @@ tl::expected WriteChunks(std::ofstream &file, } // namespace +void FileUtil::SortByCreationTime(std::vector &files) { + using TimeKey = std::tuple; + std::vector> timed_files; + timed_files.reserve(files.size()); + for (const auto &entry : files) { + bool has_birth_time = false; + int64_t seconds = 0; + uint32_t nanoseconds = 0; +#if defined(__linux__) + struct statx info{}; + has_birth_time = + statx(AT_FDCWD, entry.path().c_str(), 0, STATX_BTIME, &info) == 0 && + (info.stx_mask & STATX_BTIME) != 0; + if (has_birth_time) { + seconds = info.stx_btime.tv_sec; + nanoseconds = info.stx_btime.tv_nsec; + } +#elif defined(__APPLE__) + struct stat info{}; + has_birth_time = stat(entry.path().c_str(), &info) == 0; + if (has_birth_time) { + seconds = info.st_birthtimespec.tv_sec; + nanoseconds = info.st_birthtimespec.tv_nsec; + } +#endif + timed_files.emplace_back(TimeKey{!has_birth_time, seconds, nanoseconds}, + entry); + } + + // No filesystem operations in the comparator: a disappearing file must + // not become equivalent to every other entry or change a comparison key. + std::sort(timed_files.begin(), timed_files.end(), + [](const auto &a, const auto &b) { + if (a.first != b.first) { + return a.first < b.first; + } + return a.second.path() < b.second.path(); + }); + for (size_t i = 0; i < files.size(); ++i) { + files[i] = std::move(timed_files[i].second); + } +} + std::string FileUtil::ResolvePathFromKey(const std::string &key, const std::string &root_dir, const std::string &fsdir) { diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 130f5f0324..b296893fcd 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -136,23 +136,7 @@ tl::expected StorageBackend::Init(uint64_t quota_bytes = 0) { << e.what(); return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - std::sort(existing_files.begin(), existing_files.end(), - [](const auto& a, const auto& b) { - struct statx stx_a, stx_b; - bool success_a = (statx(AT_FDCWD, a.path().c_str(), 0, - STATX_BTIME, &stx_a) == 0); - bool success_b = (statx(AT_FDCWD, b.path().c_str(), 0, - STATX_BTIME, &stx_b) == 0); - - if (!success_a || !success_b) { - return false; - } - - if (stx_a.stx_btime.tv_sec == stx_b.stx_btime.tv_sec) { - return stx_a.stx_btime.tv_nsec < stx_b.stx_btime.tv_nsec; - } - return stx_a.stx_btime.tv_sec < stx_b.stx_btime.tv_sec; - }); + FileUtil::SortByCreationTime(existing_files); bool eviction_needed = false; { std::unique_lock space_lock(space_mutex_); diff --git a/mooncake-store/tests/file_util_test.cpp b/mooncake-store/tests/file_util_test.cpp index e53eaf5ca5..add44db19c 100644 --- a/mooncake-store/tests/file_util_test.cpp +++ b/mooncake-store/tests/file_util_test.cpp @@ -2,11 +2,18 @@ #include #include +#include +#include #include #include #include +#include +#include #include +#include +#include + #include "common/file_util.h" namespace mooncake::test { @@ -46,6 +53,125 @@ class FileUtilTest : public ::testing::Test { std::string tmp_dir_; }; +TEST_F(FileUtilTest, SortByCreationTime_EmptyAndSingle) { + std::vector files; + FileUtil::SortByCreationTime(files); + EXPECT_TRUE(files.empty()); + + const auto path = fs::path(tmp_dir()) / "one"; + std::ofstream(path) << "data"; + files.emplace_back(path); + FileUtil::SortByCreationTime(files); + ASSERT_EQ(files.size(), 1); + EXPECT_EQ(files.front().path(), path); +} + +TEST_F(FileUtilTest, SortByCreationTime_MissingFileKeepsKnownFilesOrdered) { + const auto early = fs::path(tmp_dir()) / "z_early"; + const auto late = fs::path(tmp_dir()) / "a_late"; + const auto missing = fs::path(tmp_dir()) / "removed_after_scan"; + std::ofstream(early) << "early"; + std::ofstream(missing) << "temporary"; + const fs::directory_entry missing_entry(missing); + ASSERT_TRUE(fs::remove(missing)); + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + std::ofstream(late) << "late"; + +#if defined(__linux__) + struct statx early_stat{}, late_stat{}; + ASSERT_EQ(statx(AT_FDCWD, early.c_str(), 0, STATX_BTIME, &early_stat), 0); + ASSERT_EQ(statx(AT_FDCWD, late.c_str(), 0, STATX_BTIME, &late_stat), 0); + if (!(early_stat.stx_mask & STATX_BTIME) || + !(late_stat.stx_mask & STATX_BTIME) || + std::tie(early_stat.stx_btime.tv_sec, early_stat.stx_btime.tv_nsec) >= + std::tie(late_stat.stx_btime.tv_sec, late_stat.stx_btime.tv_nsec)) { + GTEST_SKIP() << "Filesystem did not provide distinct birth times"; + } +#elif defined(__APPLE__) + struct stat early_stat{}, late_stat{}; + ASSERT_EQ(stat(early.c_str(), &early_stat), 0); + ASSERT_EQ(stat(late.c_str(), &late_stat), 0); + if (std::tie(early_stat.st_birthtimespec.tv_sec, + early_stat.st_birthtimespec.tv_nsec) >= + std::tie(late_stat.st_birthtimespec.tv_sec, + late_stat.st_birthtimespec.tv_nsec)) { + GTEST_SKIP() << "Filesystem did not provide distinct birth times"; + } +#else + GTEST_SKIP() << "Birth-time queries are not supported on this platform"; +#endif + + const std::vector entries{ + fs::directory_entry(early), fs::directory_entry(late), missing_entry}; + std::vector order{0, 1, 2}; + do { + std::vector files; + for (int index : order) files.push_back(entries[index]); + FileUtil::SortByCreationTime(files); + ASSERT_EQ(files.size(), 3); + EXPECT_EQ(files[0].path(), early); + EXPECT_EQ(files[1].path(), late); + EXPECT_EQ(files[2].path(), missing); + } while (std::next_permutation(order.begin(), order.end())); + EXPECT_TRUE(fs::exists(early)); + EXPECT_TRUE(fs::exists(late)); +} + +TEST_F(FileUtilTest, SortByCreationTime_RetainsAndOrdersFailedQueries) { + const auto a = fs::path(tmp_dir()) / "a"; + const auto z = fs::path(tmp_dir()) / "z"; + std::ofstream(a) << "a"; + std::ofstream(z) << "z"; + std::vector files{fs::directory_entry(z), + fs::directory_entry(a)}; + ASSERT_TRUE(fs::remove(a)); + ASSERT_TRUE(fs::remove(z)); + FileUtil::SortByCreationTime(files); + ASSERT_EQ(files.size(), 2); + EXPECT_EQ(files[0].path(), a); + EXPECT_EQ(files[1].path(), z); +} + +TEST_F(FileUtilTest, SortByCreationTime_EqualBirthTimesUsePath) { + const auto a = fs::path(tmp_dir()) / "a"; + const auto z = fs::path(tmp_dir()) / "z"; + std::ofstream(z) << "data"; + fs::create_hard_link(z, a); + std::vector files{fs::directory_entry(z), + fs::directory_entry(a)}; + FileUtil::SortByCreationTime(files); + ASSERT_EQ(files.size(), 2); + EXPECT_EQ(files[0].path(), a); + EXPECT_EQ(files[1].path(), z); + EXPECT_EQ(ReadFileContent(a.string()), "data"); + EXPECT_EQ(ReadFileContent(z.string()), "data"); +} + +#if defined(__linux__) +TEST_F(FileUtilTest, SortByCreationTime_RetainsReadableFileWithoutBirthTime) { + const fs::path without_birth_time = "/proc/self/status"; + struct statx info{}; + ASSERT_EQ( + statx(AT_FDCWD, without_birth_time.c_str(), 0, STATX_BTIME, &info), 0); + if (info.stx_mask & STATX_BTIME) { + GTEST_SKIP() << "procfs provides birth times on this system"; + } + const auto known = fs::path(tmp_dir()) / "known"; + std::ofstream(known) << "data"; + ASSERT_EQ(statx(AT_FDCWD, known.c_str(), 0, STATX_BTIME, &info), 0); + if (!(info.stx_mask & STATX_BTIME)) { + GTEST_SKIP() << "Temporary filesystem does not provide birth times"; + } + std::vector files{ + fs::directory_entry(without_birth_time), fs::directory_entry(known)}; + FileUtil::SortByCreationTime(files); + ASSERT_EQ(files.size(), 2); + EXPECT_EQ(files[0].path(), known); + EXPECT_EQ(files[1].path(), without_birth_time); + EXPECT_FALSE(ReadFileContent(without_birth_time.string()).empty()); +} +#endif + // ========== SaveStringToFile ========== TEST_F(FileUtilTest, SaveStringToFile_Basic) { From 137303a38ca293ac4407a88980c037b6a32fa455 Mon Sep 17 00:00:00 2001 From: Le1zyCatt <2490162471@qq.com> Date: Tue, 28 Jul 2026 11:56:51 +0800 Subject: [PATCH 13/14] [TENT] Complete native transport integration with receiver credit (cherry picked from commit d78e20dd3f9568c8223fbdbb3e98e00f6ed34c67) --- docs/source/design/tent/overview.md | 1 + .../design/tent/ub-native-validation.md | 322 ++++++++ docs/source/performance/mooncake/tebench.md | 3 +- .../tent/config/transfer-engine.json | 10 + .../tent/include/tent/common/types.h | 6 + .../tent/include/tent/rpc/rpc.h | 1 + .../include/tent/runtime/admission_queue.h | 12 + .../tent/include/tent/runtime/control_plane.h | 97 ++- .../include/tent/runtime/receiver_credit.h | 299 ++++++- .../tent/include/tent/runtime/segment.h | 13 + .../tent/runtime/transfer_engine_impl.h | 81 +- .../tent/include/tent/runtime/transport.h | 7 + .../tent/include/tent/transport/rdma/slice.h | 3 + .../tent/transport/tcp/tcp_transport.h | 11 +- .../tent/include/tent/transport/ub/endpoint.h | 6 + .../tent/include/tent/transport/ub/slice.h | 3 + .../tent/include/tent/transport/ub/workers.h | 2 + .../tent/src/runtime/admission_queue.cpp | 104 ++- .../tent/src/runtime/control_plane.cpp | 434 ++++++++-- .../tent/src/runtime/proxy_manager.cpp | 15 + .../tent/src/runtime/receiver_credit.cpp | 642 ++++++++++++--- .../tent/src/runtime/transfer_engine_impl.cpp | 740 +++++++++++++++++- .../src/transport/rdma/rdma_transport.cpp | 5 + .../tent/src/transport/rdma/workers.cpp | 18 + .../tent/src/transport/tcp/tcp_transport.cpp | 50 +- .../tent/src/transport/ub/endpoint.cpp | 24 +- .../tent/src/transport/ub/ub_transport.cpp | 108 ++- .../tent/src/transport/ub/workers.cpp | 121 ++- .../tent/tests/admission_queue_test.cpp | 210 +++++ .../tent/tests/receiver_credit_test.cpp | 358 ++++++++- .../tent/tests/request_merge_test.cpp | 1 + .../tests/runtime_queue_dispatch_test.cpp | 452 ++++++++++- .../tent/tests/transport_selector_test.cpp | 12 + .../tent/tests/ub_native_data_path_test.cpp | 55 +- 34 files changed, 3918 insertions(+), 308 deletions(-) create mode 100644 docs/source/design/tent/ub-native-validation.md diff --git a/docs/source/design/tent/overview.md b/docs/source/design/tent/overview.md index 6b37df87e7..30c5e1fdf9 100644 --- a/docs/source/design/tent/overview.md +++ b/docs/source/design/tent/overview.md @@ -139,4 +139,5 @@ failover :maxdepth: 1 testing +ub-native-validation ::: diff --git a/docs/source/design/tent/ub-native-validation.md b/docs/source/design/tent/ub-native-validation.md new file mode 100644 index 0000000000..c7294823b8 --- /dev/null +++ b/docs/source/design/tent/ub-native-validation.md @@ -0,0 +1,322 @@ +# TENT Native UB Hardware Validation Guide + +This guide validates the TENT-native UB transport on machines with a real +UMDK/URMA provider. It covers build linkage, two-node READ/WRITE correctness, +multi-rail scheduling, fault recovery, retryable shutdown, and a performance +baseline. The repository's CI uses an injected adapter; the hardware cases in +this guide are therefore required before declaring native UB production-ready. + +## 1. Record the Test Environment + +Run every case on the same Mooncake commit and record the following on both +hosts: + +```bash +git rev-parse HEAD +uname -a +gcc --version | head -1 +cmake --version | head -1 +urma_admin -l +ldconfig -p | grep liburma +``` + +The two hosts must have mutually reachable EIDs and TCP reachability for the +TENT P2P control plane. Use the vendor-supported tool to verify each active UB +port, for example: + +```bash +urma_admin -p urma0 +``` + +If UMDK is installed outside the system search paths, export its library path +before configuring and running Mooncake: + +```bash +export LD_LIBRARY_PATH=/path/to/umdk/lib:${LD_LIBRARY_PATH} +``` + +Unset `MC_TENT_CONF` for the commands below unless the referenced file +explicitly enables `transports/ub`. The environment config replaces, rather +than merges with, the configuration assembled by `tebench`. + +```bash +unset MC_TENT_CONF +``` + +## 2. Build and Verify Real URMA Linkage + +Configure a Debug build first. CMake discovers standard UMDK installations +automatically; pass `URMA_SYSTEM_INCLUDE_DIR` and `URMA_LIBRARY` only for a +non-standard installation. + +```bash +cmake -S . -B build-ub-debug \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DUSE_CUDA=OFF \ + -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BENCHMARK=ON \ + -DCMAKE_BUILD_TYPE=Debug + +cmake --build build-ub-debug -j"$(nproc)" --target \ + tent_ub_core_test \ + tent_ub_teardown_test \ + tent_ub_native_data_path_test \ + tent_transport_selector_test \ + tebench +``` + +Confirm that CMake found a real library and that `tebench` has a runtime +dependency on it: + +```bash +grep '^URMA_LIBRARY:' build-ub-debug/CMakeCache.txt +ldd build-ub-debug/mooncake-transfer-engine/benchmark/tebench | grep urma +``` + +Both commands must show the intended UMDK installation. A header-only fallback +is sufficient for mock tests but **not** for hardware validation. + +Run the deterministic mock suite before touching the devices: + +```bash +ctest --test-dir build-ub-debug --output-on-failure \ + -R 'tent_(ub_core|ub_teardown|ub_native_data_path|transport_selector)_test' +``` + +Expected result: all selected tests pass. In particular, the teardown test +must demonstrate that a failed unregister/delete retains ownership and that a +second cleanup call releases the same handle. It also verifies that a failed +device unpublishes every old endpoint and cannot recover until every JFC has +reported healthy in the current failure epoch. + +## 3. Two-Node Correctness + +Use `HOST_A` as the target and `HOST_B` as the initiator. The examples use a +fixed control-plane port so peer-restart tests can reuse the same segment name. + +On `HOST_A`: + +```bash +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --rpc_server_port=18001 \ + --xport_type=ub \ + --seg_type=DRAM \ + --total_buffer_size=1073741824 2>&1 | tee target-ub.log +``` + +Copy the printed segment name. On `HOST_B`, run a consistency-checked mixed +transfer first: + +```bash +TARGET_SEG='' + +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --target_seg_name="${TARGET_SEG}" \ + --xport_type=ub \ + --tent_transport_hint=ub \ + --seg_type=DRAM \ + --op_type=mix \ + --check_consistency=true \ + --total_buffer_size=1073741824 \ + --start_block_size=4096 \ + --max_block_size=1048576 \ + --start_batch_size=1 \ + --max_batch_size=8 \ + --start_num_threads=1 \ + --max_num_threads=4 \ + --duration=10 2>&1 | tee initiator-mix.log +``` + +Then repeat with `--op_type=read` and `--op_type=write`. All rows must finish, +consistency checking must report no mismatch, and neither host may log a +fallback to a non-UB transport. This covers both upward-facing TENT APIs and +native URMA READ/WRITE without converting requests to the Classic TE path. + +## 4. Multi-Device and Pressure Scheduling + +Leave `transports/ub/device_filter` empty so every discovered active EID is +eligible. Confirm that the startup topology lists at least two `ub::eid` +entries on each host, then run a fixed high-concurrency case: + +```bash +build-ub-debug/mooncake-transfer-engine/benchmark/tebench \ + --backend=tent \ + --metadata_type=p2p \ + --target_seg_name="${TARGET_SEG}" \ + --xport_type=ub \ + --tent_transport_hint=ub \ + --seg_type=DRAM \ + --op_type=write \ + --total_buffer_size=1073741824 \ + --start_block_size=65536 \ + --max_block_size=65536 \ + --start_batch_size=16 \ + --max_batch_size=16 \ + --start_num_threads=16 \ + --max_num_threads=16 \ + --duration=60 2>&1 | tee initiator-multirail.log +``` + +Monitor per-device traffic with the platform's supported counters. Every +healthy eligible rail should make progress; saturating one path must not stop +posting on another path with free quota. Record aggregate bandwidth and the +per-device byte deltas. + +## 5. Fault and Recovery Cases + +Run these cases with a long-lived embedding harness that keeps the initiator's +TENT engine alive across each fault. Preserve both host logs, note exact fault +timestamps, and record endpoint generations before and after recovery. A new +`tebench` process may be used for basic reconnect smoke testing, but it is not +sufficient evidence for generation isolation. + +### Peer restart and endpoint rebuild + +1. Stop `HOST_A` with `SIGINT` while transfers are active. +2. Start it again with the same `--rpc_server_port=18001` command. +3. In the still-running `HOST_B` harness, discover the republished segment and + submit READ and WRITE again without reinstalling the local TENT transport. + +The old endpoint generation must not be reused. The live engine must establish +a new generation and complete new transfers; stale completions must not change +a new task. Save both generation values as part of the PASS evidence. + +### Remote segment republish + +Using the embedding application's TENT API, unregister and re-register the +target buffer at the same virtual address while the engines remain alive. +Publish the returned `BufferDesc`, then submit another READ and WRITE. The new +metadata generation must be imported and the stale imported segment must be +released. Do not use `tebench` for this case because it registers its buffers +only once at process startup. + +### UB port down/up + +1. Disable one UB port using the deployment's vendor-approved administration + procedure; do not power-cycle a host as the first test. +2. Confirm that traffic continues on another healthy rail. +3. Re-enable the port and wait at least the configured endpoint cooldown. +4. Confirm that the recovered rail resumes traffic and that a new endpoint is + built before posting. + +The current recovery loop first unpublishes every endpoint backed by the failed +device. It reactivates the context only after every provider JFC has polled +successfully in the latest failure epoch, the cooldown has elapsed, and all old +WRs have drained. This supports a transient port/link failure when the provider +keeps Context/JFC handles valid. A destructive device reset that invalidates +those handles still requires transport reinstall and must be reported as +unsupported for in-place recovery. + +### Cancellation and timeout + +Exercise cancellation through the embedding application's TENT API while a +large multi-slice transfer is active. For timeout testing, drop the UB path +without stopping the control-plane TCP connection. The transfer must either +complete on another rail or reach a terminal timeout/failure; it must not hang +and shutdown must still drain safely. + +### Transport fallback + +The `tebench --xport_type=ub` commands above deliberately disable every other +transport, so they validate UB path failover but **cannot** validate transport +fallback. Use an embedding test harness with both UB and TCP (or RDMA) enabled, +and a selector policy that prefers UB but permits the second transport. Submit +the request with `transport_hint=UNSPEC`; do not pin it to UB. + +1. Verify the first consistency-checked READ and WRITE select UB. +2. Disable every usable UB path while requests remain active. +3. Verify the runtime reaches a terminal result for the failed UB attempt and + resubmits through TCP/RDMA rather than leaving the task pending. +4. Restore UB, wait for its cooldown, and verify later policy-selected traffic + can use UB again. + +Capture selector/failover logs or transport metrics that identify both the +failed UB attempt and the replacement transport. A successful copy without +this evidence does not prove transport fallback. + +## 6. Repeated Shutdown and Sanitizers + +Configure a separate sanitizer build: + +```bash +cmake -S . -B build-ub-asan \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DUSE_CUDA=OFF \ + -DBUILD_UNIT_TESTS=ON \ + -DBUILD_BENCHMARK=ON \ + -DCMAKE_BUILD_TYPE=Debug \ + -DCMAKE_CXX_FLAGS='-fsanitize=address,undefined -fno-omit-frame-pointer' \ + -DCMAKE_EXE_LINKER_FLAGS='-fsanitize=address,undefined' \ + -DCMAKE_SHARED_LINKER_FLAGS='-fsanitize=address,undefined' + +cmake --build build-ub-asan -j"$(nproc)" --target tebench +export ASAN_OPTIONS=detect_leaks=1:halt_on_error=1 +export UBSAN_OPTIONS=halt_on_error=1:print_stacktrace=1 +``` + +Use an embedding test harness to perform at least 100 +install-register-transfer-unregister-uninstall cycles **in one process**. +Include ten cycles where a UB port or peer is unavailable during uninstall. If +unregister or uninstall returns a provider `busy`/error, keep the engine alive +and retry the same cleanup operation; the retry must act on the retained native +handle and eventually succeed after the fault is removed. Also run 100 +target/initiator start-transfer-`SIGINT` process cycles as a supplementary +check. + +A cycle is successful only if shutdown returns, the next cycle can reuse the +same buffers and control-plane port, and ASan/UBSan reports no UAF, double free, +or leaked native ownership graph. Record the first and retry return codes plus +the provider log around every injected cleanup failure. + +## 7. Performance Baseline + +Use a `RelWithDebInfo` build and keep CPU affinity, NUMA placement, block size, +batch size, concurrency, UMDK version, and device set fixed. Record: + +| Case | Required measurements | +| --- | --- | +| 4 KiB, batch 1, one thread | average, P99, and P999 latency | +| 64 KiB and 1 MiB, one thread | single-stream bandwidth and CPU usage | +| 64 KiB and 1 MiB, 16 threads | aggregate bandwidth and CPU usage | +| Multi-device, 16 threads | aggregate and per-device bandwidth | +| First transfer to a peer | endpoint connection time | + +Run the same matrix against Classic TE UB using the existing Classic UB +benchmark on the same hosts. Any material regression needs a profile and an +explanation; do not compare results from different UMDK versions or topology. + +## 8. Result Checklist + +Attach the following to the review: + +```text +Mooncake commit: +Host/OS/kernel: +UMDK package and liburma path: +Devices/EIDs: +Debug USE_TENT+USE_UB build: PASS/FAIL +Mock UB tests: PASS/FAIL +Two-node READ: PASS/FAIL +Two-node WRITE: PASS/FAIL +Consistency-checked mix: PASS/FAIL +Multi-device pressure scheduling: PASS/FAIL +Peer restart / endpoint rebuild: PASS/FAIL +Segment republish: PASS/FAIL +Port down/up: PASS/FAIL +Cancellation and timeout: PASS/FAIL +UB to TCP/RDMA transport fallback: PASS/FAIL +100-cycle ASan/UBSan shutdown: PASS/FAIL +Classic UB baseline artifact: +TENT UB baseline artifact: +Known provider errors or unsupported cases: +``` + +Do not mark the hardware gate complete from bandwidth alone. Correctness, +generation isolation, retryable resource release, and sanitizer-clean shutdown +are independent acceptance requirements. diff --git a/docs/source/performance/mooncake/tebench.md b/docs/source/performance/mooncake/tebench.md index 47180d8abf..e39cec0fc9 100644 --- a/docs/source/performance/mooncake/tebench.md +++ b/docs/source/performance/mooncake/tebench.md @@ -397,7 +397,8 @@ gpu_id + thread_id **Transport (TENT only)** -* `--xport_type` : `rdma | shm | mnnvl | gds | iouring`. Selects a single +* `--xport_type` : `rdma | tcp | shm | mnnvl | nvlink | gds | iouring | ub | + sunrise_link`. Selects a single transport to enable (all others are disabled). Empty means no transport is explicitly enabled or disabled by tebench β€” the engine reads the transport enable list from the `MC_TENT_CONF` config file (see Section diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json index 01880b843e..2ba9fcd185 100644 --- a/mooncake-transfer-engine/tent/config/transfer-engine.json +++ b/mooncake-transfer-engine/tent/config/transfer-engine.json @@ -12,6 +12,16 @@ "max_failover_attempts": 3, "enable_auto_failover_on_poll": true, "enable_progress_worker": false, + "receiver_credit": { + "enable": false, + "data_bytes": 1073741824, + "request_slots": 1024, + "freshness_ttl_ms": 1000, + "max_freshness_ttl_ms": 60000, + "max_senders": 1024, + "max_authority_entries": 4096, + "max_sender_entries": 4096 + }, "metrics": { "enabled": true, "http_port": 9100, diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index b9d758180d..64562bf9d1 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -177,6 +177,12 @@ struct Request { // and docs/source/design/tent/deadline-scheduling.md. uint64_t deadline_ns = 0; IntentType intent_type = IntentType::INTENT_UNSPEC; + // Internal runtime fence for staging sub-transfers that are covered by + // their parent request's receiver-credit reservation. Public callers + // leave these fields zero. + uint64_t receiver_credit_session_high = 0; + uint64_t receiver_credit_session_low = 0; + uint64_t receiver_credit_epoch = 0; }; enum TransferStatusEnum { diff --git a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index daeac4916d..a19f0224ff 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -49,6 +49,7 @@ enum RpcFuncID { NotifySegmentUpdated, // Appended to preserve the numeric values of the existing RPCs. BootstrapUb, + ExchangeReceiverCredit, }; class ClientPool; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index 5623e9194e..786686988e 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -147,6 +147,12 @@ class LocalTransferAdmissionQueue { size_t max_owners, size_t max_bytes, std::vector* dropped_owner_ids = nullptr); + // Return an owner that could not be handed off to its selected transport + // from the dispatch window to the queue. Only a Dispatching owner can be + // deferred. Outstanding owner/byte accounting remains charged while the + // owner waits for a later dispatch attempt. + Status defer(QueueOwnerId owner_id); + // Install the step-3 degradation policy inputs. Optional; without it the // queue never drops (default behavior). now defaults to steady_clock. void setDegradationPolicy(BandwidthProvider bandwidth_provider, @@ -192,10 +198,16 @@ class LocalTransferAdmissionQueue { QueueOwnerKind kind{QueueOwnerKind::User}; TransportType transport{UNSPEC}; bool degradation_eligible{false}; + // EDF owners that yield because a downstream resource is unavailable + // move to the next scheduling round. This preserves progress without + // relying on a rotated deque remaining globally deadline-sorted. + uint64_t scheduling_round{0}; QueueState state{QueueState::Queued}; TransferStatusEnum terminal_status{TransferStatusEnum::PENDING}; }; + void enqueueOwner(QueueOwnerId owner_id); + QueueLimits limits_; Status limits_status_; QueueOwnerId next_owner_id_{1}; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h index ab1ce5f3e1..93b439a996 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -77,25 +78,32 @@ struct UbBootstrapDesc { std::vector jetty_uasids; uint64_t endpoint_generation = 0; uint64_t segment_generation = 0; + uint64_t receiver_credit_session_high = 0; + uint64_t receiver_credit_session_low = 0; + uint64_t receiver_credit_epoch = 0; std::vector capabilities; std::string reply_msg; }; inline void to_json(nlohmann::json& j, const UbBootstrapDesc& desc) { - j = nlohmann::json{{"protocol_version", desc.protocol_version}, - {"segment_name", desc.segment_name}, - {"local_nic_path", desc.local_nic_path}, - {"peer_nic_path", desc.peer_nic_path}, - {"local_device_name", desc.local_device_name}, - {"local_device_id", desc.local_device_id}, - {"local_eid_index", desc.local_eid_index}, - {"local_eid", desc.local_eid}, - {"jetty_ids", desc.jetty_ids}, - {"jetty_uasids", desc.jetty_uasids}, - {"endpoint_generation", desc.endpoint_generation}, - {"segment_generation", desc.segment_generation}, - {"capabilities", desc.capabilities}, - {"reply_msg", desc.reply_msg}}; + j = nlohmann::json{ + {"protocol_version", desc.protocol_version}, + {"segment_name", desc.segment_name}, + {"local_nic_path", desc.local_nic_path}, + {"peer_nic_path", desc.peer_nic_path}, + {"local_device_name", desc.local_device_name}, + {"local_device_id", desc.local_device_id}, + {"local_eid_index", desc.local_eid_index}, + {"local_eid", desc.local_eid}, + {"jetty_ids", desc.jetty_ids}, + {"jetty_uasids", desc.jetty_uasids}, + {"endpoint_generation", desc.endpoint_generation}, + {"segment_generation", desc.segment_generation}, + {"receiver_credit_session_high", desc.receiver_credit_session_high}, + {"receiver_credit_session_low", desc.receiver_credit_session_low}, + {"receiver_credit_epoch", desc.receiver_credit_epoch}, + {"capabilities", desc.capabilities}, + {"reply_msg", desc.reply_msg}}; } inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { @@ -114,14 +122,44 @@ inline void from_json(const nlohmann::json& j, UbBootstrapDesc& desc) { desc.jetty_uasids = j.value("jetty_uasids", std::vector{}); desc.endpoint_generation = j.value("endpoint_generation", uint64_t{0}); desc.segment_generation = j.value("segment_generation", uint64_t{0}); + desc.receiver_credit_session_high = + j.value("receiver_credit_session_high", uint64_t{0}); + desc.receiver_credit_session_low = + j.value("receiver_credit_session_low", uint64_t{0}); + desc.receiver_credit_epoch = j.value("receiver_credit_epoch", uint64_t{0}); desc.capabilities = j.value("capabilities", std::vector{}); desc.reply_msg = j.value("reply_msg", ""); } +// SendData/RecvData v2 keeps the legacy address/length prefix, then adds a +// magic discriminator and the receiver-credit session fence. A new receiver +// can still parse legacy non-credit requests, while credit-required receivers +// reject any request without an exact current fence. +constexpr uint64_t kXferDataProtocolMagic = 0x54454e5458465232ULL; +constexpr uint64_t kXferDataV2LegacyRejectAddress = + std::numeric_limits::max(); +struct LegacyXferDataDesc { + uint64_t peer_mem_addr; + uint64_t length; +}; struct XferDataDesc { + // An old receiver interprets the first two fields as the legacy header. + // The impossible address makes it reject v2 instead of copying the + // appended fence bytes into user memory. + uint64_t legacy_reject_addr; + uint64_t length; + uint64_t protocol_magic; uint64_t peer_mem_addr; - size_t length; + uint64_t receiver_credit_session_high; + uint64_t receiver_credit_session_low; + uint64_t receiver_credit_epoch; }; +struct XferDataReplyHeader { + uint64_t protocol_magic; +}; +static_assert(sizeof(LegacyXferDataDesc) == 16); +static_assert(sizeof(XferDataDesc) == 56); +static_assert(sizeof(XferDataReplyHeader) == 8); // Standard TCP RPC payload boundaries. READ retains the server's 1 GiB cap; // WRITE uses a 32-bit RPC attachment that includes XferDataDesc. @@ -139,6 +177,10 @@ using OnReceiveBootstrap = using OnReceiveUbBootstrap = std::function; +using OnExchangeReceiverCredit = + std::function; + using OnNotify = std::function; class ControlClient { @@ -164,13 +206,24 @@ class ControlClient { const UbBootstrapDesc& request, UbBootstrapDesc& response); + static Status exchangeReceiverCredit( + const std::string& server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& response); + static Status sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length); + size_t length, + uint64_t receiver_credit_session_high = 0, + uint64_t receiver_credit_session_low = 0, + uint64_t receiver_credit_epoch = 0); static Status recvData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length); + size_t length, + uint64_t receiver_credit_session_high = 0, + uint64_t receiver_credit_session_low = 0, + uint64_t receiver_credit_epoch = 0); static Status notify(const std::string& server_addr, const Notification& message); @@ -224,6 +277,11 @@ class ControlService { ub_bootstrap_callback_ = callback; } + void setReceiverCreditCallback(const OnExchangeReceiverCredit& callback) { + std::lock_guard lock(receiver_credit_callback_mutex_); + receiver_credit_callback_ = callback; + } + void setNotifyCallback(const OnNotify& callback); Status start(uint16_t& port, bool ipv6_ = false, size_t threads = 1); @@ -237,6 +295,9 @@ class ControlService { void onBootstrapUb(const std::string_view& request, std::string& response); + void onExchangeReceiverCredit(const std::string_view& request, + std::string& response); + void onSendData(const std::string_view& request, std::string& response); void onRecvData(const std::string_view& request, std::string& response); @@ -280,6 +341,8 @@ class ControlService { std::mutex notify_cb_mutex_; std::condition_variable notify_cb_cv_; size_t notify_callbacks_in_flight_ = 0; + std::mutex receiver_credit_callback_mutex_; + OnExchangeReceiverCredit receiver_credit_callback_; OnNotify notify_callback_; static thread_local const ControlService* active_notify_service_; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h index 09b243ab13..76ad717327 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h @@ -6,93 +6,352 @@ #include #include +#include #include +#include #include #include #include #include "tent/common/status.h" +#include "tent/thirdparty/nlohmann/json.h" namespace mooncake::tent { +// A receiver session changes whenever a receiver process restarts. A sender +// instance has the same representation, but remains a distinct type so the +// two identities cannot be accidentally exchanged at an API boundary. struct ReceiverSessionId { uint64_t high{0}, low{0}; - bool operator==(const ReceiverSessionId& o) const { - return high == o.high && low == o.low; - } + bool operator==(const ReceiverSessionId&) const = default; + [[nodiscard]] bool empty() const noexcept { return high == 0 && low == 0; } }; + +struct SenderInstanceId { + uint64_t high{0}, low{0}; + bool operator==(const SenderInstanceId&) const = default; + [[nodiscard]] bool empty() const noexcept { return high == 0 && low == 0; } +}; + struct CreditKey { ReceiverSessionId receiver_session; - uint64_t sender_peer{0}; + SenderInstanceId sender_instance; uint32_t qos_class{0}; - bool operator==(const CreditKey& o) const { - return receiver_session == o.receiver_session && - sender_peer == o.sender_peer && qos_class == o.qos_class; - } + bool operator==(const CreditKey&) const = default; }; + struct CreditKeyHash { size_t operator()(const CreditKey&) const noexcept; }; + enum class CreditResource : uint16_t { DataBytes = 1, RequestSlots, StagingSlots, ConsumerSlots }; + constexpr size_t kCreditResourceCount = 4; +constexpr uint16_t kReceiverCreditProtocolVersion = 1; +constexpr uint16_t kReceiverCreditRequired = 1U << 0; + +// A cumulative grant. grant_total never decreases within a session/epoch. struct CreditAmount { - CreditResource resource; + CreditResource resource{CreditResource::DataBytes}; uint64_t grant_total{0}; }; + +// A cumulative counter used for release reports and desired grant totals. +struct CreditCounter { + CreditResource resource{CreditResource::DataBytes}; + uint64_t total{0}; +}; + struct CreditCharge { std::vector> resources; }; + +struct ReceiverCreditAdvertV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}; + uint16_t flags{kReceiverCreditRequired}; + ReceiverSessionId receiver_session_id; + uint64_t epoch{0}; + uint32_t freshness_ttl_ms{0}; + std::vector resources; + // Explicit capacities let a sender reject a request that can never fit + // instead of leaving it indefinitely deferred behind the credit gate. + std::vector capacities; +}; + struct ReceiverCreditUpdateV1 { - uint16_t schema_version{1}, flags{0}; + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; uint32_t qos_class{0}; ReceiverSessionId receiver_session_id; uint64_t epoch{0}, sequence{0}; uint32_t freshness_ttl_ms{0}; std::vector grants; }; + +// One cumulative exchange is sufficient for initial credit, replenishment, +// replay recovery, and terminal release reporting. desired_grant_totals is +// absolute (not a delta), making retries idempotent. +struct ReceiverCreditExchangeRequestV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; + ReceiverSessionId expected_receiver_session_id; + uint64_t epoch{0}; + SenderInstanceId sender_instance_id; + uint32_t qos_class{0}; + uint64_t known_update_sequence{0}; + uint64_t report_sequence{0}; + std::vector released_totals; + std::vector desired_grant_totals; +}; + +struct ReceiverCreditExchangeReplyV1 { + uint16_t schema_version{kReceiverCreditProtocolVersion}, flags{0}; + SenderInstanceId sender_instance_id; + ReceiverCreditUpdateV1 update; + std::string reply_msg; +}; + enum class CreditUpdateDisposition : uint8_t { Applied, DuplicateOrOld, SequenceGap }; -// Private, sender-side state model. It has no network or Admission integration. +using CreditNowProvider = std::function; + +// Sender-side cumulative ledger. A successful transport handoff consumes a +// grant permanently; terminal completion reports a cumulative release to the +// receiver, which is the only authority allowed to mint a replacement grant. class SenderCreditLedger { public: - explicit SenderCreditLedger(size_t max_entries = 1024) - : max_entries_(max_entries) {} + explicit SenderCreditLedger(size_t max_entries = 1024, + uint32_t max_freshness_ttl_ms = 60000, + CreditNowProvider now_provider = {}); + Status activate(const CreditKey&, uint64_t epoch); - // Removes an exactly matched epoch after the caller has fenced and drained - // (or failed) its transport-owned work. An old cleanup cannot erase a - // reactivated, newer epoch. Status deactivate(const CreditKey&, uint64_t epoch); + Status deactivateSession(const ReceiverSessionId&); Status applyUpdate(const CreditKey&, const ReceiverCreditUpdateV1&, CreditUpdateDisposition&); + Status tryReserve(const CreditKey&, const CreditCharge&); - // Only for work not yet handed to a transport; completions need a new - // grant. + // Only for a failed submit that is proven to have handed no work to a + // transport. Consumption remains cumulative and the same amount is added + // to released, so the receiver (rather than this sender) reclaims it. Status rollbackReservation(const CreditKey&, const CreditCharge&); + // Records terminal completion of committed work. It deliberately does + // not reduce consumed; the receiver must authorize replacement credit. + Status releaseCommitted(const CreditKey&, const CreditCharge&); + + // Builds the next idempotent exchange. The desired totals request enough + // credit for `next_charge` in addition to all locally consumed credit. + Status prepareExchange(const CreditKey&, const CreditCharge& next_charge, + ReceiverCreditExchangeRequestV1&); + // Reports terminal releases without asking the receiver to replenish an + // idle sender. A later prepareExchange() requests credit on demand. + Status prepareReleaseExchange(const CreditKey&, + ReceiverCreditExchangeRequestV1&); + Status available(const CreditKey&, CreditResource, uint64_t&) const; Status consumed(const CreditKey&, CreditResource, uint64_t&) const; + Status released(const CreditKey&, CreditResource, uint64_t&) const; private: struct Entry { - uint64_t epoch{0}, last_sequence{0}; + uint64_t epoch{0}, last_sequence{0}, report_sequence{0}; + uint64_t expires_at_ns{0}; bool has_update{false}; - std::array grants{}, consumed{}; + std::array grants{}, consumed{}, + released{}; }; + static Status resourceIndex(CreditResource, size_t&); static Status normalize(const CreditCharge&, std::array&); + uint64_t nowNanos() const; + Status requireFresh(const Entry&) const; + mutable std::mutex mutex_; const size_t max_entries_; + const uint32_t max_freshness_ttl_ms_; + CreditNowProvider now_provider_; + std::unordered_map entries_; +}; + +struct ReceiverCreditLimits { + std::array capacity{}; + uint32_t freshness_ttl_ms{1000}; + size_t max_senders{1024}; + // Replay history is retained for the receiver session so an old initial + // demand cannot be accepted again after all of its grants are released. + // The bound is fail-closed; entries are reset only by a new receiver + // session, which also fences every old request on the wire. + size_t max_entries{4096}; +}; + +// Receiver-side grant authority. All resources in one desired request are +// granted atomically, and aggregate outstanding authorization can never exceed +// the configured capacity. Expiry is sender-side freshness only: this class +// conservatively does not recycle an unreachable sender without a transport +// fence proving its old DMA is gone. +class ReceiverCreditAuthority { + public: + ReceiverCreditAuthority(ReceiverSessionId session, uint64_t epoch, + ReceiverCreditLimits limits); + + Status status() const { return status_; } + ReceiverCreditAdvertV1 advert() const; + Status exchange(const ReceiverCreditExchangeRequestV1&, + ReceiverCreditExchangeReplyV1&); + Status outstanding(CreditResource, uint64_t&) const; + + private: + struct Entry { + uint64_t last_report_sequence{0}; + uint64_t update_sequence{0}; + std::array grants{}, released{}; + }; + + static Status validateLimits(const ReceiverCreditLimits&); + static Status normalizeCounters(const std::vector&, + std::array&, + std::array&); + void fillReply(const CreditKey&, const Entry&, + ReceiverCreditExchangeReplyV1&) const; + + ReceiverSessionId session_; + uint64_t epoch_{0}; + ReceiverCreditLimits limits_; + Status status_; + mutable std::mutex mutex_; + std::array outstanding_{}; std::unordered_map entries_; }; +inline void to_json(nlohmann::json& j, const ReceiverSessionId& id) { + j = nlohmann::json{{"high", id.high}, {"low", id.low}}; +} +inline void from_json(const nlohmann::json& j, ReceiverSessionId& id) { + j.at("high").get_to(id.high); + j.at("low").get_to(id.low); +} +inline void to_json(nlohmann::json& j, const SenderInstanceId& id) { + j = nlohmann::json{{"high", id.high}, {"low", id.low}}; +} +inline void from_json(const nlohmann::json& j, SenderInstanceId& id) { + j.at("high").get_to(id.high); + j.at("low").get_to(id.low); +} +inline void to_json(nlohmann::json& j, const CreditResource& resource) { + j = static_cast(resource); +} +inline void from_json(const nlohmann::json& j, CreditResource& resource) { + resource = static_cast(j.get()); +} +inline void to_json(nlohmann::json& j, const CreditAmount& amount) { + j = nlohmann::json{{"resource", amount.resource}, + {"grant_total", amount.grant_total}}; +} +inline void from_json(const nlohmann::json& j, CreditAmount& amount) { + j.at("resource").get_to(amount.resource); + j.at("grant_total").get_to(amount.grant_total); +} +inline void to_json(nlohmann::json& j, const CreditCounter& counter) { + j = nlohmann::json{{"resource", counter.resource}, + {"total", counter.total}}; +} +inline void from_json(const nlohmann::json& j, CreditCounter& counter) { + j.at("resource").get_to(counter.resource); + j.at("total").get_to(counter.total); +} +inline void to_json(nlohmann::json& j, const ReceiverCreditAdvertV1& advert) { + j = nlohmann::json{{"schema_version", advert.schema_version}, + {"flags", advert.flags}, + {"receiver_session_id", advert.receiver_session_id}, + {"epoch", advert.epoch}, + {"freshness_ttl_ms", advert.freshness_ttl_ms}, + {"resources", advert.resources}, + {"capacities", advert.capacities}}; +} +inline void from_json(const nlohmann::json& j, ReceiverCreditAdvertV1& advert) { + advert.schema_version = j.value("schema_version", uint16_t{0}); + advert.flags = j.value("flags", uint16_t{0}); + j.at("receiver_session_id").get_to(advert.receiver_session_id); + advert.epoch = j.value("epoch", uint64_t{0}); + advert.freshness_ttl_ms = j.value("freshness_ttl_ms", uint32_t{0}); + advert.resources = j.value("resources", std::vector{}); + advert.capacities = j.value("capacities", std::vector{}); +} +inline void to_json(nlohmann::json& j, const ReceiverCreditUpdateV1& update) { + j = nlohmann::json{{"schema_version", update.schema_version}, + {"flags", update.flags}, + {"qos_class", update.qos_class}, + {"receiver_session_id", update.receiver_session_id}, + {"epoch", update.epoch}, + {"sequence", update.sequence}, + {"freshness_ttl_ms", update.freshness_ttl_ms}, + {"grants", update.grants}}; +} +inline void from_json(const nlohmann::json& j, ReceiverCreditUpdateV1& update) { + update.schema_version = j.value("schema_version", uint16_t{0}); + update.flags = j.value("flags", uint16_t{0}); + update.qos_class = j.value("qos_class", uint32_t{0}); + j.at("receiver_session_id").get_to(update.receiver_session_id); + update.epoch = j.value("epoch", uint64_t{0}); + update.sequence = j.value("sequence", uint64_t{0}); + update.freshness_ttl_ms = j.value("freshness_ttl_ms", uint32_t{0}); + update.grants = j.value("grants", std::vector{}); +} +inline void to_json(nlohmann::json& j, + const ReceiverCreditExchangeRequestV1& request) { + j = nlohmann::json{ + {"schema_version", request.schema_version}, + {"flags", request.flags}, + {"expected_receiver_session_id", request.expected_receiver_session_id}, + {"epoch", request.epoch}, + {"sender_instance_id", request.sender_instance_id}, + {"qos_class", request.qos_class}, + {"known_update_sequence", request.known_update_sequence}, + {"report_sequence", request.report_sequence}, + {"released_totals", request.released_totals}, + {"desired_grant_totals", request.desired_grant_totals}}; +} +inline void from_json(const nlohmann::json& j, + ReceiverCreditExchangeRequestV1& request) { + request.schema_version = j.value("schema_version", uint16_t{0}); + request.flags = j.value("flags", uint16_t{0}); + j.at("expected_receiver_session_id") + .get_to(request.expected_receiver_session_id); + request.epoch = j.value("epoch", uint64_t{0}); + j.at("sender_instance_id").get_to(request.sender_instance_id); + request.qos_class = j.value("qos_class", uint32_t{0}); + request.known_update_sequence = + j.value("known_update_sequence", uint64_t{0}); + request.report_sequence = j.value("report_sequence", uint64_t{0}); + request.released_totals = + j.value("released_totals", std::vector{}); + request.desired_grant_totals = + j.value("desired_grant_totals", std::vector{}); +} +inline void to_json(nlohmann::json& j, + const ReceiverCreditExchangeReplyV1& reply) { + j = nlohmann::json{{"schema_version", reply.schema_version}, + {"flags", reply.flags}, + {"sender_instance_id", reply.sender_instance_id}, + {"update", reply.update}, + {"reply_msg", reply.reply_msg}}; +} +inline void from_json(const nlohmann::json& j, + ReceiverCreditExchangeReplyV1& reply) { + reply.schema_version = j.value("schema_version", uint16_t{0}); + reply.flags = j.value("flags", uint16_t{0}); + j.at("sender_instance_id").get_to(reply.sender_instance_id); + j.at("update").get_to(reply.update); + reply.reply_msg = j.value("reply_msg", ""); +} + } // namespace mooncake::tent #endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h index 63d73df312..1f6d8f846a 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,7 @@ #include "tent/common/concurrent/rw_spinlock.h" #include "tent/common/concurrent/thread_local_storage.h" #include "tent/runtime/metastore.h" +#include "tent/runtime/receiver_credit.h" #include "tent/rpc/rpc.h" #include "tent/runtime/topology.h" #include "tent/thirdparty/nlohmann/json.h" @@ -107,6 +109,11 @@ struct MemorySegmentDesc { // int maps to TransportType enum std::unordered_map transport_attrs; + // Optional runtime-level receiver-credit advertisement. Per-sender + // grants are never published in SegmentDesc; they are obtained through + // the credit exchange RPC. + std::optional receiver_credit; + public: // Get transport attributes for a specific transport type std::string& getTransportAttrs(int transport_type) { @@ -129,6 +136,7 @@ inline void to_json(json& j, const MemorySegmentDesc& m) { {"devices", m.devices}, {"topology", m.topology.toString()}, {"transport_attrs", m.transport_attrs}}; + if (m.receiver_credit) j["receiver_credit"] = *m.receiver_credit; } inline void from_json(const json& j, MemorySegmentDesc& m) { @@ -142,6 +150,11 @@ inline void from_json(const json& j, MemorySegmentDesc& m) { if (j.contains("transport_attrs")) { j.at("transport_attrs").get_to(m.transport_attrs); } + m.receiver_credit.reset(); + if (j.contains("receiver_credit")) { + m.receiver_credit = + j.at("receiver_credit").get(); + } } struct FileSegmentDesc { diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 6251927bda..3e4ac91f2c 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -33,6 +33,7 @@ #include "tent/common/status.h" #include "tent/common/types.h" #include "tent/runtime/admission_queue.h" +#include "tent/runtime/receiver_credit.h" #include "tent/runtime/transport.h" #include "tent/runtime/transport_selector.h" #include "tent/runtime/hp_tcp_transport_config.h" @@ -463,7 +464,7 @@ class TransferEngineImpl { QueueOwnerKind owner_kind) const; Status prepareSubmit(Batch* batch, const std::vector& request_list, - PreparedSubmit& prepared, + QueueOwnerKind owner_kind, PreparedSubmit& prepared, bool require_post_submit_cancellation = false); Status commitPreparedSubmit(Batch* batch, const PreparedSubmit& prepared); @@ -482,6 +483,43 @@ class TransferEngineImpl { Status dispatchQueuedOwner(QueueOwnerId owner_id); + struct ReceiverCreditTarget; + struct ReceiverCreditReservation; + + Status resolveReceiverCreditTarget(const Request& request, + ReceiverCreditTarget& target) const; + + Status acquireReceiverCredit(const Request& request, + ReceiverCreditReservation& reservation, + bool& acquired); + + Status validateReceiverCreditCapacity( + const Request& request, const ReceiverCreditTarget& target) const; + + void bindReceiverCreditFence(Transport::SubBatchRef, + const ReceiverCreditReservation&) const; + + void bindReceiverCreditFence(Transport::SubBatchRef, const Request&) const; + + void rollbackReceiverCredit(ReceiverCreditReservation& reservation); + + void releaseReceiverCredit(ReceiverCreditReservation& reservation); + + void scheduleReceiverCreditRelease(const CreditKey&, SegmentID, + const std::string&, uint64_t epoch); + + bool tryReportReceiverCreditRelease(const CreditKey&); + + void progressReceiverCreditReleases(); + + void observeReceiverCreditSession(SegmentID, + const ReceiverSessionId& session); + + Status exchangeReceiverCredit( + SegmentID target_id, const std::string& rpc_server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply); + Status markQueuedOwnerSubmitted(QueueOwnerId owner_id); Status finishQueuedOwner(QueueOwnerId owner_id, @@ -578,12 +616,46 @@ class TransferEngineImpl { std::chrono::microseconds progress_fallback_interval{50000}; }; + struct ReceiverCreditConfig { + bool sender_enabled{false}; + bool receiver_enabled{false}; + ReceiverCreditLimits limits{}; + size_t max_sender_entries{4096}; + uint32_t max_freshness_ttl_ms{60000}; + }; + + struct ReceiverCreditTarget { + bool required{false}; + ReceiverCreditAdvertV1 advert{}; + std::string rpc_server_addr; + }; + + struct ReceiverCreditReservation { + bool active{false}; + bool committed{false}; + CreditKey key{}; + CreditCharge charge{}; + SegmentID target_id{LOCAL_SEGMENT_ID}; + std::string rpc_server_addr; + uint64_t epoch{0}; + }; + + struct PendingReceiverCreditRelease { + SegmentID target_id{LOCAL_SEGMENT_ID}; + std::string rpc_server_addr; + uint64_t epoch{0}; + uint32_t attempts{0}; + std::chrono::steady_clock::time_point retry_at{}; + }; + struct QueuedOwnerState { Batch* batch{nullptr}; size_t owner_task_id{0}; std::vector public_task_ids; size_t byte_charge{0}; bool in_dispatch_window{false}; + bool receiver_credit_exempt{false}; + ReceiverCreditReservation credit; }; private: @@ -617,7 +689,14 @@ class TransferEngineImpl { bool notify_rpc_fallback_{true}; bool enable_progress_worker_{false}; RuntimeQueueConfig runtime_queue_config_; + ReceiverCreditConfig receiver_credit_config_; std::unique_ptr runtime_queue_; + SenderInstanceId sender_instance_id_; + std::unique_ptr sender_credit_ledger_; + std::unique_ptr receiver_credit_authority_; + std::unordered_map + pending_credit_releases_; + std::unordered_map receiver_credit_sessions_; std::unordered_map queued_owners_; size_t dispatch_inflight_owners_{0}; size_t dispatch_inflight_bytes_{0}; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index 0299475a7a..00bef306af 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -59,6 +59,13 @@ class Transport { // no pool (default spray). Carried like device_mask, from the matched // SelectionPolicy down to each RdmaTask. std::string qp_pool; + // Runtime receiver-credit authorization bound to the next submitted + // task. Native transports copy this fence into their asynchronous + // work so a refreshed segment cannot post under a different receiver + // session than the one that granted capacity. + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; BatchID progress_batch_id{0}; std::function notify_progress; }; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index de1f7a558b..1c97e6167d 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -107,6 +107,9 @@ struct RdmaTask { // no pool selected: slices spray across all data QPs as before. Resolved // from SelectionResult.qp_pool at task creation. std::string qp_pool; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; std::atomic success_slices{0}; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h index 1f0fc77d9a..d6ed97e340 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h @@ -55,6 +55,9 @@ struct TcpTask { std::atomic transferred_bytes{0}; bool non_replayable_failure{false}; // Published by status_word. uint64_t target_addr = 0; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; TcpTask() = default; TcpTask(TcpTask &&other) noexcept @@ -65,7 +68,10 @@ struct TcpTask { transferred_bytes( other.transferred_bytes.load(std::memory_order_relaxed)), non_replayable_failure(other.non_replayable_failure), - target_addr(other.target_addr) {} + target_addr(other.target_addr), + receiver_credit_session_high(other.receiver_credit_session_high), + receiver_credit_session_low(other.receiver_credit_session_low), + receiver_credit_epoch(other.receiver_credit_epoch) {} TcpTask(const TcpTask &) = delete; TcpTask &operator=(const TcpTask &) = delete; }; @@ -122,8 +128,7 @@ class TcpTransport : public Transport { Status doTransferWithRetry(TcpTask *task); - Status findRemoteSegment(uint64_t dest_addr, uint64_t length, - uint64_t target_id, std::string &rpc_server_addr); + Status findRemoteSegment(const TcpTask &task, std::string &rpc_server_addr); private: bool installed_; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h index 3dfe4533ba..3679bffaf3 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/endpoint.h @@ -10,6 +10,7 @@ #include #include #include +#include #include #include "tent/common/status.h" @@ -116,6 +117,11 @@ class UbEndpoint final : public std::enable_shared_from_this { [[nodiscard]] uint64_t peerGeneration() const noexcept { return peer_generation_.load(std::memory_order_acquire); } + [[nodiscard]] bool peerSupportsCapability( + std::string_view capability) const; + [[nodiscard]] bool peerMatchesReceiverCredit(uint64_t session_high, + uint64_t session_low, + uint64_t epoch) const; [[nodiscard]] uint64_t outstandingWrs() const noexcept { return outstanding_wrs_.load(std::memory_order_relaxed); } diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h index 3b70c3dc21..7428c5b0d8 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/slice.h @@ -147,6 +147,9 @@ struct UbSliceSpec { size_t length{0}; size_t request_offset{0}; uint32_t max_retries{0}; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; }; struct UbAttemptToken { diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h index a2e9031fe8..7f3fc4e1dd 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h @@ -147,6 +147,8 @@ class UbWorkers final { Status start(); Status stop(); + Status submitBatch(const std::vector& tasks, + uint64_t device_mask = ~0ULL); Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL); Status cancel(const UbTask::Ptr& task); diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 74c19416ad..3a273bf953 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -78,6 +78,30 @@ Status validateLimits(const QueueLimits& limits) { LocalTransferAdmissionQueue::LocalTransferAdmissionQueue(QueueLimits limits) : limits_(limits), limits_status_(validateLimits(limits)) {} +void LocalTransferAdmissionQueue::enqueueOwner(QueueOwnerId owner_id) { + fifo_.push_back(owner_id); + if (!limits_.deadline_aware) return; + + // Promotion and defer can both disturb the deque's prior global order. + // Re-establish the explicit (round, EDF) ordering on admission instead of + // passing a potentially rotated range to an ordered insertion algorithm. + // stable_sort preserves FIFO order for equal deadlines. + std::stable_sort( + fifo_.begin(), fifo_.end(), [this](QueueOwnerId lhs, QueueOwnerId rhs) { + const auto lhs_it = owners_.find(lhs); + const auto rhs_it = owners_.find(rhs); + if (lhs_it == owners_.end()) return false; + if (rhs_it == owners_.end()) return true; + const auto& lhs_owner = lhs_it->second; + const auto& rhs_owner = rhs_it->second; + if (lhs_owner.scheduling_round != rhs_owner.scheduling_round) { + return lhs_owner.scheduling_round < rhs_owner.scheduling_round; + } + return deadlineKey(lhs_owner.request.deadline_ns) < + deadlineKey(rhs_owner.request.deadline_ns); + }); +} + Status LocalTransferAdmissionQueue::tryAdmit( const QueueSubmit& submit, std::vector& admitted_owner_ids) { admitted_owner_ids.clear(); @@ -181,6 +205,17 @@ Status LocalTransferAdmissionQueue::tryAdmit( owner.kind = owner_input.kind; owner.transport = owner_input.transport; owner.degradation_eligible = owner_input.degradation_eligible; + if (limits_.deadline_aware) { + for (const auto queued_id : fifo_) { + auto queued_it = owners_.find(queued_id); + if (queued_it != owners_.end() && + queued_it->second.state == QueueState::Queued) { + owner.scheduling_round = + std::max(owner.scheduling_round, + queued_it->second.scheduling_round); + } + } + } owners_.emplace(owner_id, owner); public_to_owner_[{submit.batch_token, owner_input.owner_task_id}] = @@ -189,26 +224,8 @@ Status LocalTransferAdmissionQueue::tryAdmit( public_to_owner_[{submit.batch_token, derived_task_id}] = owner_id; } // RFC #2519 step 2: keep fifo_ ordered on admission so pickForDispatch - // never has to re-sort. Default (deadline_aware == false) appends in - // strict FIFO. When deadline-aware, insert at the earliest-deadline- - // first position; upper_bound places a new owner *after* existing - // owners with the same deadline, preserving FIFO order among ties. - if (limits_.deadline_aware) { - const uint64_t key = deadlineKey(owner.request.deadline_ns); - auto pos = std::upper_bound( - fifo_.begin(), fifo_.end(), key, - [this](uint64_t k, QueueOwnerId id) { - auto it = owners_.find(id); - uint64_t d = - (it == owners_.end()) - ? std::numeric_limits::max() - : deadlineKey(it->second.request.deadline_ns); - return k < d; - }); - fifo_.insert(pos, owner_id); - } else { - fifo_.push_back(owner_id); - } + // never has to re-sort. + enqueueOwner(owner_id); admitted_owner_ids.push_back(owner_id); } @@ -271,16 +288,26 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( .count())) : 0; - // Deadline proximity promotion: partition fifo_ so owners with critical - // slack (deadline approaching within promotion_slack_ns) appear before - // owners with comfortable slack or no deadline. stable_partition preserves - // relative EDF order within each group. + // Deadline proximity promotion applies only inside the earliest + // scheduling round. A resource-starved owner that defer() moved to the + // next round must not jump back ahead of current-round waiters merely + // because its deadline is urgent. if (promotion_enabled) { + uint64_t current_round = std::numeric_limits::max(); + for (const auto owner_id : fifo_) { + auto owner_it = owners_.find(owner_id); + if (owner_it != owners_.end() && + owner_it->second.state == QueueState::Queued) { + current_round = + std::min(current_round, owner_it->second.scheduling_round); + } + } std::stable_partition(fifo_.begin(), fifo_.end(), [&](QueueOwnerId id) { auto it = owners_.find(id); if (it == owners_.end() || it->second.state != QueueState::Queued) { return false; } + if (it->second.scheduling_round != current_round) return false; const uint64_t dl = it->second.request.deadline_ns; if (dl == 0 || dl <= now_ns) return false; return (dl - now_ns) < limits_.promotion_slack_ns; @@ -383,6 +410,35 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( return picked; } +Status LocalTransferAdmissionQueue::defer(QueueOwnerId owner_id) { + if (owner_id == 0) { + return Status::InvalidArgument("invalid queue owner id" LOC_MARK); + } + + auto owner_it = owners_.find(owner_id); + if (owner_it == owners_.end()) { + return Status::InvalidEntry("queue owner not found" LOC_MARK); + } + if (owner_it->second.state != QueueState::Dispatching) { + return Status::InvalidEntry("queue owner is not dispatching" LOC_MARK); + } + + // A resource-deferred EDF owner advances one scheduling round. It yields + // to every owner already waiting in this round while retaining EDF order + // among owners that yielded into the next one. + if (limits_.deadline_aware) { + if (owner_it->second.scheduling_round == + std::numeric_limits::max()) { + return Status::InternalError( + "admission scheduling round exhausted" LOC_MARK); + } + ++owner_it->second.scheduling_round; + } + owner_it->second.state = QueueState::Queued; + enqueueOwner(owner_id); + return Status::OK(); +} + Status LocalTransferAdmissionQueue::complete( QueueOwnerId owner_id, TransferStatusEnum terminal_status) { if (owner_id == 0) { diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index 0d1624f19b..71cf343bf6 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -52,6 +53,113 @@ thread_local const ControlService* ControlService::active_notify_service_ = nullptr; thread_local CoroRpcAgent tl_rpc_agent; +namespace { + +constexpr uint64_t kDelegatedCreditProtocolMagic = 0x54454e54444c4732ULL; + +struct ParsedXferData { + uint64_t peer_mem_addr{0}; + size_t length{0}; + size_t header_size{0}; + bool version_two{false}; + uint64_t receiver_credit_session_high{0}; + uint64_t receiver_credit_session_low{0}; + uint64_t receiver_credit_epoch{0}; +}; + +Status parseXferData(const std::string_view& request, bool carries_payload, + ParsedXferData& parsed) { + if (request.size() < sizeof(LegacyXferDataDesc)) { + return Status::InvalidArgument("request too short" LOC_MARK); + } + LegacyXferDataDesc legacy; + memcpy(&legacy, request.data(), sizeof(legacy)); + const uint64_t length = le64toh(legacy.length); + if (length > std::numeric_limits::max()) { + return Status::InvalidArgument("transfer length overflow" LOC_MARK); + } + + parsed = ParsedXferData{}; + parsed.peer_mem_addr = le64toh(legacy.peer_mem_addr); + parsed.length = static_cast(length); + parsed.header_size = sizeof(LegacyXferDataDesc); + + if (request.size() >= sizeof(XferDataDesc)) { + XferDataDesc current; + memcpy(¤t, request.data(), sizeof(current)); + const bool magic_matches = + le64toh(current.protocol_magic) == kXferDataProtocolMagic; + const bool size_matches = + carries_payload + ? parsed.length <= std::numeric_limits::max() - + sizeof(XferDataDesc) && + request.size() == sizeof(XferDataDesc) + parsed.length + : request.size() == sizeof(XferDataDesc); + if (magic_matches && size_matches) { + parsed.header_size = sizeof(XferDataDesc); + parsed.version_two = true; + parsed.peer_mem_addr = le64toh(current.peer_mem_addr); + parsed.receiver_credit_session_high = + le64toh(current.receiver_credit_session_high); + parsed.receiver_credit_session_low = + le64toh(current.receiver_credit_session_low); + parsed.receiver_credit_epoch = + le64toh(current.receiver_credit_epoch); + } + } + + if (carries_payload) { + if (parsed.length > + std::numeric_limits::max() - parsed.header_size || + request.size() != parsed.header_size + parsed.length) { + return Status::InvalidArgument( + "invalid transfer request size" LOC_MARK); + } + } else if (request.size() != parsed.header_size) { + return Status::InvalidArgument( + "invalid transfer request header size" LOC_MARK); + } + return Status::OK(); +} + +Status validateXferReceiverCreditFence(const SegmentDescRef& local_desc, + const ParsedXferData& request) { + const bool has_fence = request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + if (!local_desc || local_desc->type != SegmentType::Memory) { + return has_fence ? Status::InvalidEntry( + "receiver-credit target is not memory" LOC_MARK) + : Status::OK(); + } + + const auto& advert = local_desc->getMemory().receiver_credit; + if (!advert) { + return has_fence + ? Status::InvalidEntry( + "receiver-credit session is unavailable" LOC_MARK) + : Status::OK(); + } + if (advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.empty() || advert->epoch == 0) { + return Status::InvalidEntry( + "invalid local receiver-credit advertisement" LOC_MARK); + } + if (!request.version_two || + request.receiver_credit_session_high != + advert->receiver_session_id.high || + request.receiver_credit_session_low != + advert->receiver_session_id.low || + request.receiver_credit_epoch != advert->epoch) { + return Status::InvalidEntry( + "receiver-credit session fence mismatch" LOC_MARK); + } + return Status::OK(); +} + +} // namespace + Status ControlClient::getSegmentDesc(const std::string& server_addr, std::string& response) { std::string request; @@ -113,29 +221,83 @@ Status ControlClient::bootstrapUb(const std::string& server_addr, return Status::OK(); } +Status ControlClient::exchangeReceiverCredit( + const std::string& server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& response) { + std::string request_raw, response_raw; + json encoded = request; + request_raw = encoded.dump(); + CHECK_STATUS(tl_rpc_agent.call(server_addr, ExchangeReceiverCredit, + request_raw, response_raw)); + try { + response = + json::parse(response_raw).get(); + } catch (const std::exception& error) { + return Status::MalformedJson( + std::string("Malformed receiver-credit response: ") + error.what() + + LOC_MARK); + } + if (response.schema_version != kReceiverCreditProtocolVersion || + response.flags != 0) { + return Status::RpcServiceError( + "Unsupported receiver-credit response version" LOC_MARK); + } + if (!response.reply_msg.empty()) { + return Status::RpcServiceError(response.reply_msg); + } + return Status::OK(); +} + Status ControlClient::sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length) { - // Check before addition/allocation or reading the caller's source buffer. + size_t length, + uint64_t receiver_credit_session_high, + uint64_t receiver_credit_session_low, + uint64_t receiver_credit_epoch) { if (length > kTcpMaxWriteBytes) return Status::InvalidArgument("TCP WRITE exceeds RPC payload limit"); - std::string response; - XferDataDesc desc{htole64(peer_mem_addr), htole64(length)}; - std::string request; + std::string request, response; + const bool has_any_fence = receiver_credit_session_high != 0 || + receiver_credit_session_low != 0 || + receiver_credit_epoch != 0; + const bool has_complete_fence = (receiver_credit_session_high != 0 || + receiver_credit_session_low != 0) && + receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial SendData receiver-credit fence" LOC_MARK); + } + if (length > std::numeric_limits::max() - sizeof(XferDataDesc)) { + return Status::InvalidArgument( + "SendData request size overflow" LOC_MARK); + } + size_t header_size = sizeof(LegacyXferDataDesc); request.reserve(sizeof(XferDataDesc) + length); - request.resize(sizeof(XferDataDesc)); - memcpy(request.data(), &desc, sizeof(desc)); + if (has_complete_fence) { + const XferDataDesc desc{htole64(kXferDataV2LegacyRejectAddress), + htole64(static_cast(length)), + htole64(kXferDataProtocolMagic), + htole64(peer_mem_addr), + htole64(receiver_credit_session_high), + htole64(receiver_credit_session_low), + htole64(receiver_credit_epoch)}; + header_size = sizeof(desc); + request.resize(header_size); + memcpy(request.data(), &desc, sizeof(desc)); + } else { + const LegacyXferDataDesc desc{htole64(peer_mem_addr), + htole64(static_cast(length))}; + request.resize(header_size); + memcpy(request.data(), &desc, sizeof(desc)); + } auto& loader = Platform::getLoader(); if (loader.getMemoryType(local_mem_addr) == MTYPE_CPU) { - // Single copy into the RPC attachment; avoids the resize() zero-fill - // and the extra copy in call(). request.append(reinterpret_cast(local_mem_addr), length); } else { - // resize() zero-fills the payload, so an unchecked copy failure would - // ship zeros that the peer stores successfully and reports COMPLETED. - request.resize(sizeof(XferDataDesc) + length); + request.resize(header_size + length); CHECK_STATUS( - loader.copy(request.data() + sizeof(desc), local_mem_addr, length)); + loader.copy(request.data() + header_size, local_mem_addr, length)); } auto status = tl_rpc_agent.callOwned(server_addr, SendData, std::move(request), response); @@ -146,19 +308,63 @@ Status ControlClient::sendData(const std::string& server_addr, Status ControlClient::recvData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, - size_t length) { + size_t length, + uint64_t receiver_credit_session_high, + uint64_t receiver_credit_session_low, + uint64_t receiver_credit_epoch) { if (length > kTcpMaxReadBytes) return Status::InvalidArgument("TCP READ exceeds RPC payload limit"); std::string request, response; - XferDataDesc desc{htole64(peer_mem_addr), htole64(length)}; - request.resize(sizeof(XferDataDesc)); - memcpy(&request[0], &desc, sizeof(desc)); + const bool has_any_fence = receiver_credit_session_high != 0 || + receiver_credit_session_low != 0 || + receiver_credit_epoch != 0; + const bool has_complete_fence = (receiver_credit_session_high != 0 || + receiver_credit_session_low != 0) && + receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial RecvData receiver-credit fence" LOC_MARK); + } + if (has_complete_fence) { + const XferDataDesc desc{htole64(kXferDataV2LegacyRejectAddress), + htole64(static_cast(length)), + htole64(kXferDataProtocolMagic), + htole64(peer_mem_addr), + htole64(receiver_credit_session_high), + htole64(receiver_credit_session_low), + htole64(receiver_credit_epoch)}; + request.resize(sizeof(desc)); + memcpy(request.data(), &desc, sizeof(desc)); + } else { + const LegacyXferDataDesc desc{htole64(peer_mem_addr), + htole64(static_cast(length))}; + request.resize(sizeof(desc)); + memcpy(request.data(), &desc, sizeof(desc)); + } auto status = tl_rpc_agent.call(server_addr, RecvData, request, response); if (!status.ok()) return status; - if (response.size() != length) + if (!has_complete_fence) { + if (response.size() != length) { + return Status::RpcServiceError( + "RecvData failed: invalid or rejected legacy response"); + } + return Platform::getLoader().copy(local_mem_addr, response.data(), + length); + } + if (length > + std::numeric_limits::max() - sizeof(XferDataReplyHeader) || + response.size() != sizeof(XferDataReplyHeader) + length) { + return Status::RpcServiceError( + "RecvData failed: invalid or rejected response"); + } + XferDataReplyHeader header; + memcpy(&header, response.data(), sizeof(header)); + if (le64toh(header.protocol_magic) != kXferDataProtocolMagic) { return Status::RpcServiceError( - response.empty() ? "RecvData failed: empty response" : response); - return Platform::getLoader().copy(local_mem_addr, response.data(), length); + "RecvData failed: unsupported response version"); + } + return Platform::getLoader().copy( + local_mem_addr, response.data() + sizeof(XferDataReplyHeader), length); } inline void to_json(nlohmann::json& j, const Notification& n) { @@ -184,11 +390,26 @@ Status ControlClient::probe(const std::string& server_addr) { } inline void to_json(json& j, const Request& r) { - j = json{{"opcode", r.opcode == Request::READ ? "READ" : "WRITE"}, - {"source", reinterpret_cast(r.source)}, - {"target_id", r.target_id}, - {"target_offset", r.target_offset}, - {"length", r.length}}; + const bool has_complete_fence = (r.receiver_credit_session_high != 0 || + r.receiver_credit_session_low != 0) && + r.receiver_credit_epoch != 0; + j = json{ + {"opcode", r.opcode == Request::READ ? "READ" : "WRITE"}, + {"source", reinterpret_cast(r.source)}, + {"target_id", has_complete_fence ? std::numeric_limits::max() + : r.target_id}, + {"target_offset", has_complete_fence + ? std::numeric_limits::max() + : r.target_offset}, + {"length", r.length}, + {"receiver_credit_session_high", r.receiver_credit_session_high}, + {"receiver_credit_session_low", r.receiver_credit_session_low}, + {"receiver_credit_epoch", r.receiver_credit_epoch}}; + if (has_complete_fence) { + j["delegate_protocol_magic"] = kDelegatedCreditProtocolMagic; + j["actual_target_id"] = r.target_id; + j["actual_target_offset"] = r.target_offset; + } } inline void from_json(const json& j, Request& r) { @@ -201,13 +422,35 @@ inline void from_json(const json& j, Request& r) { throw std::runtime_error("Invalid opcode"); r.source = reinterpret_cast(j.at("source").get()); - r.target_id = j.at("target_id").get(); - r.target_offset = j.at("target_offset").get(); + const bool fenced_delegate = + j.value("delegate_protocol_magic", uint64_t{0}) == + kDelegatedCreditProtocolMagic; + r.target_id = fenced_delegate ? j.at("actual_target_id").get() + : j.at("target_id").get(); + r.target_offset = fenced_delegate + ? j.at("actual_target_offset").get() + : j.at("target_offset").get(); r.length = j.at("length").get(); + r.receiver_credit_session_high = + j.value("receiver_credit_session_high", uint64_t{0}); + r.receiver_credit_session_low = + j.value("receiver_credit_session_low", uint64_t{0}); + r.receiver_credit_epoch = j.value("receiver_credit_epoch", uint64_t{0}); } Status ControlClient::delegate(const std::string& server_addr, const Request& request) { + const bool has_any_fence = request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + const bool has_complete_fence = + (request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0) && + request.receiver_credit_epoch != 0; + if (has_any_fence && !has_complete_fence) { + return Status::InvalidArgument( + "partial delegated receiver-credit fence" LOC_MARK); + } std::string request_raw, response_raw; json j = request; request_raw = j.dump(); @@ -309,6 +552,11 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onBootstrapUb(request, response); }); + rpc_server_->registerFunction( + ExchangeReceiverCredit, + [this](const std::string_view& request, std::string& response) { + onExchangeReceiverCredit(request, response); + }); rpc_server_->registerFunction( SendData, [this](const std::string_view& request, std::string& response) { @@ -364,6 +612,10 @@ ControlService::~ControlService() { std::lock_guard lock(ub_bootstrap_callback_mutex_); ub_bootstrap_callback_ = {}; } + { + std::lock_guard lock(receiver_credit_callback_mutex_); + receiver_credit_callback_ = {}; + } rpc_server_.reset(); } @@ -507,37 +759,62 @@ void ControlService::onBootstrapUb(const std::string_view& request, response = j.dump(); } +void ControlService::onExchangeReceiverCredit(const std::string_view& request, + std::string& response) { + ReceiverCreditExchangeReplyV1 reply; + try { + auto parsed = json::parse(std::string(request)) + .get(); + Status status = Status::InvalidEntry( + "receiver credit callback is not registered" LOC_MARK); + { + // Serialize callback replacement with invocation so teardown + // cannot destroy the authority while an exchange is in flight. + std::lock_guard lock(receiver_credit_callback_mutex_); + if (receiver_credit_callback_) { + status = receiver_credit_callback_(parsed, reply); + } + } + if (!status.ok()) reply.reply_msg = status.ToString(); + } catch (const std::exception& error) { + reply.reply_msg = + std::string("Malformed receiver-credit request: ") + error.what(); + } + json encoded = reply; + response = encoded.dump(); +} + void ControlService::onSendData(const std::string_view& request, std::string& response) { - if (request.size() < sizeof(XferDataDesc)) { - response = "SendData failed: request too short"; + ParsedXferData parsed; + auto status = parseXferData(request, true, parsed); + if (!status.ok()) { + response = "SendData failed: " + status.ToString(); return; } - XferDataDesc* desc = (XferDataDesc*)request.data(); auto local_desc = manager_->getLocal(); - auto peer_mem_addr = le64toh(desc->peer_mem_addr); - auto length = le64toh(desc->length); + status = validateXferReceiverCreditFence(local_desc, parsed); + if (!status.ok()) { + response = "SendData failed: " + status.ToString(); + return; + } - // Validate request size to prevent buffer over-read - if (length > kTcpMaxWriteBytes || - length > request.size() - sizeof(XferDataDesc)) { + if (parsed.length > kTcpMaxWriteBytes) { response = "SendData failed: invalid request size"; return; } - - if (auto* buffer = local_desc->findBuffer(peer_mem_addr, length)) { + if (auto* buffer = + local_desc->findBuffer(parsed.peer_mem_addr, parsed.length)) { if (buffer->permission != kGlobalReadWrite) { response = "SendData failed: remote write permission denied"; return; } - auto status = - Platform::getLoader().copy((void*)peer_mem_addr, &desc[1], length); - if (!status.ok()) { - // A non-empty response is interpreted as an RPC error by the - // client (see ControlClient::sendData). Without this the sender's - // transfer would be reported COMPLETED even though the destination - // buffer was never written. - response = "SendData failed: copy: " + status.ToString(); + auto copy_status = Platform::getLoader().copy( + reinterpret_cast(parsed.peer_mem_addr), + const_cast(request.data() + parsed.header_size), + parsed.length); + if (!copy_status.ok()) { + response = "SendData failed: copy: " + copy_status.ToString(); } } else { response = "SendData failed: target address not in registered buffer"; @@ -546,46 +823,59 @@ void ControlService::onSendData(const std::string_view& request, void ControlService::onRecvData(const std::string_view& request, std::string& response) { - if (request.size() < sizeof(XferDataDesc)) { - response = "RecvData failed: request too short"; + ParsedXferData parsed; + auto status = parseXferData(request, false, parsed); + // Legacy readers accept any response of exactly the requested length. + // Fenced readers also require the reply magic; pad errors to keep both + // response formats unambiguous, including parse and session failures. + const size_t reply_header_size = + parsed.version_two ? sizeof(XferDataReplyHeader) : 0; + auto fail = [&response, &parsed, reply_header_size](std::string message) { + response = std::move(message); + if (parsed.length <= + std::numeric_limits::max() - reply_header_size && + response.size() == reply_header_size + parsed.length) { + response.push_back(' '); + } + }; + if (!status.ok()) { + fail("RecvData failed: " + status.ToString()); return; } - XferDataDesc* desc = (XferDataDesc*)request.data(); auto local_desc = manager_->getLocal(); - auto peer_mem_addr = le64toh(desc->peer_mem_addr); - auto length = le64toh(desc->length); - - // The client accepts any response of exactly `length` bytes as payload (see - // ControlClient::recvData), so an error of that size must be padded or it - // would be copied into the caller's buffer and reported as success. - auto fail = [&response, length](std::string message) { - response = std::move(message); - if (response.size() == length) response.push_back(' '); - }; - - // Validate length to prevent DoS via excessive memory allocation - if (length > kTcpMaxReadBytes) { + status = validateXferReceiverCreditFence(local_desc, parsed); + if (!status.ok()) { + fail("RecvData failed: " + status.ToString()); + return; + } + if (parsed.length > kTcpMaxReadBytes) { fail("RecvData failed: length exceeds maximum allowed"); return; } - - if (auto* buffer = local_desc->findBuffer(peer_mem_addr, length)) { + if (auto* buffer = + local_desc->findBuffer(parsed.peer_mem_addr, parsed.length)) { if (buffer->permission != kGlobalReadOnly && buffer->permission != kGlobalReadWrite) { fail("RecvData failed: remote read permission denied"); return; } auto& loader = Platform::getLoader(); - if (loader.getMemoryType((void*)peer_mem_addr) == MTYPE_CPU) { - // assign() skips the resize() zero-fill pass. - response.assign(reinterpret_cast(peer_mem_addr), - length); + auto* source = reinterpret_cast(parsed.peer_mem_addr); + response.clear(); + if (parsed.version_two) { + const XferDataReplyHeader reply{htole64(kXferDataProtocolMagic)}; + response.assign(reinterpret_cast(&reply), + sizeof(reply)); + } + if (loader.getMemoryType(source) == MTYPE_CPU) { + response.append(reinterpret_cast(source), + parsed.length); } else { - response.resize(length); - auto status = - loader.copy(response.data(), (void*)peer_mem_addr, length); - if (!status.ok()) { - fail("RecvData failed: copy: " + status.ToString()); + response.resize(reply_header_size + parsed.length); + auto copy_status = loader.copy(response.data() + reply_header_size, + source, parsed.length); + if (!copy_status.ok()) { + fail("RecvData failed: copy: " + copy_status.ToString()); } } } else { diff --git a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp index 545eca753d..3936f37361 100644 --- a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp @@ -233,6 +233,11 @@ BatchID ProxyManager::submitCrossStage(const Request& request, inter_stage.length = chunk_length; inter_stage.target_id = request.target_id; inter_stage.target_offset = remote_stage_buffer; + inter_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + inter_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + inter_stage.receiver_credit_epoch = request.receiver_credit_epoch; auto batch = impl_->allocateBatch(1); auto status = impl_->submitStagingTransfer(batch, {inter_stage}); if (!status.ok()) { @@ -283,6 +288,11 @@ Status ProxyManager::waitRemoteStage(const std::string& server_addr, remote_stage.length = chunk_length; remote_stage.target_id = LOCAL_SEGMENT_ID; remote_stage.target_offset = request.target_offset + offset; + remote_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + remote_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + remote_stage.receiver_credit_epoch = request.receiver_credit_epoch; return ControlClient::delegate(server_addr, remote_stage); } @@ -296,6 +306,11 @@ void ProxyManager::submitRemoteStage( remote_stage.length = chunk_length; remote_stage.target_id = LOCAL_SEGMENT_ID; remote_stage.target_offset = request.target_offset + offset; + remote_stage.receiver_credit_session_high = + request.receiver_credit_session_high; + remote_stage.receiver_credit_session_low = + request.receiver_credit_session_low; + remote_stage.receiver_credit_epoch = request.receiver_credit_epoch; operation = std::make_shared( server_addr, remote_stage_buffer, [server_addr, remote_stage_buffer] { ControlClient::unpinStageBufferAsync( diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp index 5ae6c14a2e..00bf199c6e 100644 --- a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp @@ -3,176 +3,622 @@ #include "tent/runtime/receiver_credit.h" +#include +#include +#include +#include + namespace mooncake::tent { +namespace { + +uint64_t steadyNowNs() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} -size_t CreditKeyHash::operator()(const CreditKey& k) const noexcept { - size_t h = std::hash{}(k.receiver_session.high); - auto mix = [&h](uint64_t v) { - h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + - (h >> 2); - }; - mix(k.receiver_session.low); - mix(k.sender_peer); - mix(k.qos_class); - return h; -} - -Status SenderCreditLedger::resourceIndex(CreditResource r, size_t& i) { - auto raw = static_cast(r); - if (raw < 1 || raw > kCreditResourceCount) +Status indexFor(CreditResource resource, size_t& index) { + const auto raw = static_cast(resource); + if (raw < 1 || raw > kCreditResourceCount) { return Status::InvalidArgument("unknown credit resource" LOC_MARK); - i = raw - 1; + } + index = raw - 1; return Status::OK(); } +void mixHash(size_t& hash, uint64_t value) { + hash ^= std::hash{}(value) + 0x9e3779b97f4a7c15ULL + (hash << 6) + + (hash >> 2); +} + +struct SenderInstanceIdHash { + size_t operator()(const SenderInstanceId& id) const noexcept { + size_t hash = std::hash{}(id.high); + mixHash(hash, id.low); + return hash; + } +}; + +} // namespace + +size_t CreditKeyHash::operator()(const CreditKey& key) const noexcept { + size_t hash = std::hash{}(key.receiver_session.high); + mixHash(hash, key.receiver_session.low); + mixHash(hash, key.sender_instance.high); + mixHash(hash, key.sender_instance.low); + mixHash(hash, key.qos_class); + return hash; +} + +SenderCreditLedger::SenderCreditLedger(size_t max_entries, + uint32_t max_freshness_ttl_ms, + CreditNowProvider now_provider) + : max_entries_(max_entries), + max_freshness_ttl_ms_(max_freshness_ttl_ms), + now_provider_(std::move(now_provider)) {} + +Status SenderCreditLedger::resourceIndex(CreditResource resource, + size_t& index) { + return indexFor(resource, index); +} + Status SenderCreditLedger::normalize( - const CreditCharge& c, std::array& out) { - out.fill(0); - if (c.resources.empty()) + const CreditCharge& charge, + std::array& normalized) { + normalized.fill(0); + if (charge.resources.empty()) { return Status::InvalidArgument("empty credit charge" LOC_MARK); - for (auto [r, amount] : c.resources) { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); - if (!amount || out[i]) + } + for (const auto& [resource, amount] : charge.resources) { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); + if (amount == 0 || normalized[index] != 0) { return Status::InvalidArgument( "zero or duplicate credit charge" LOC_MARK); - out[i] = amount; + } + normalized[index] = amount; } return Status::OK(); } -Status SenderCreditLedger::activate(const CreditKey& k, uint64_t epoch) { - if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); +uint64_t SenderCreditLedger::nowNanos() const { + return now_provider_ ? now_provider_() : steadyNowNs(); +} + +Status SenderCreditLedger::requireFresh(const Entry& entry) const { + if (!entry.has_update || entry.expires_at_ns == 0 || + nowNanos() >= entry.expires_at_ns) { + return Status::InvalidEntry("receiver credit is stale" LOC_MARK); + } + return Status::OK(); +} + +Status SenderCreditLedger::activate(const CreditKey& key, uint64_t epoch) { + if (key.receiver_session.empty() || key.sender_instance.empty() || + epoch == 0) { + return Status::InvalidArgument( + "invalid credit activation identity" LOC_MARK); + } std::lock_guard lock(mutex_); - auto existing = entries_.find(k); + auto existing = entries_.find(key); if (existing != entries_.end()) { - if (epoch < existing->second.epoch) + if (epoch < existing->second.epoch) { return Status::InvalidEntry("stale credit activation" LOC_MARK); + } if (epoch == existing->second.epoch) return Status::OK(); Entry replacement; replacement.epoch = epoch; existing->second = replacement; return Status::OK(); - } else if (entries_.size() >= max_entries_) { + } + if (entries_.size() >= max_entries_) { return Status::TooManyRequests("credit ledger entry limit" LOC_MARK); } - Entry e; - e.epoch = epoch; - entries_.emplace(k, e); + Entry entry; + entry.epoch = epoch; + entries_.emplace(key, entry); return Status::OK(); } -Status SenderCreditLedger::deactivate(const CreditKey& k, uint64_t epoch) { - if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); +Status SenderCreditLedger::deactivate(const CreditKey& key, uint64_t epoch) { + if (epoch == 0) { + return Status::InvalidArgument("zero credit epoch" LOC_MARK); + } std::lock_guard lock(mutex_); - auto existing = entries_.find(k); - if (existing == entries_.end()) return Status::OK(); // idempotent cleanup - if (existing->second.epoch != epoch) + auto existing = entries_.find(key); + if (existing == entries_.end()) return Status::OK(); + if (existing->second.epoch != epoch) { return Status::InvalidEntry("credit cleanup epoch mismatch" LOC_MARK); + } entries_.erase(existing); return Status::OK(); } -Status SenderCreditLedger::applyUpdate(const CreditKey& k, - const ReceiverCreditUpdateV1& u, +Status SenderCreditLedger::deactivateSession(const ReceiverSessionId& session) { + if (session.empty()) { + return Status::InvalidArgument( + "empty receiver credit session cleanup" LOC_MARK); + } + std::lock_guard lock(mutex_); + for (auto it = entries_.begin(); it != entries_.end();) { + if (it->first.receiver_session == session) { + it = entries_.erase(it); + } else { + ++it; + } + } + return Status::OK(); +} + +Status SenderCreditLedger::applyUpdate(const CreditKey& key, + const ReceiverCreditUpdateV1& update, CreditUpdateDisposition& disposition) { - if (u.schema_version != 1 || !u.epoch || !u.sequence) + if (update.schema_version != kReceiverCreditProtocolVersion || + update.flags != 0 || update.epoch == 0 || update.sequence == 0 || + update.freshness_ttl_ms == 0 || + update.freshness_ttl_ms > max_freshness_ttl_ms_) { return Status::InvalidArgument("invalid credit update header" LOC_MARK); - if (!(u.receiver_session_id == k.receiver_session) || - u.qos_class != k.qos_class || u.grants.size() > kCreditResourceCount) + } + if (!(update.receiver_session_id == key.receiver_session) || + update.qos_class != key.qos_class || + update.grants.size() > kCreditResourceCount) { return Status::InvalidArgument("credit update identity/size" LOC_MARK); + } + std::array proposed{}; std::array present{}; - for (auto a : u.grants) { - size_t i = 0; - CHECK_STATUS(resourceIndex(a.resource, i)); - if (present[i]) + for (const auto& amount : update.grants) { + size_t index = 0; + CHECK_STATUS(resourceIndex(amount.resource, index)); + if (present[index]) { return Status::InvalidArgument("duplicate grant resource" LOC_MARK); - present[i] = true; - proposed[i] = a.grant_total; + } + present[index] = true; + proposed[index] = amount.grant_total; } + std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || it->second.epoch != u.epoch) + auto it = entries_.find(key); + if (it == entries_.end() || it->second.epoch != update.epoch) { return Status::InvalidEntry("inactive or stale credit epoch" LOC_MARK); - auto& e = it->second; - if (e.has_update && u.sequence <= e.last_sequence) { + } + auto& entry = it->second; + if (entry.has_update && update.sequence <= entry.last_sequence) { disposition = CreditUpdateDisposition::DuplicateOrOld; return Status::OK(); } - // `grants` is a partial cumulative update: each resource present in this - // message replaces that resource's cumulative grant total, while omitted - // resources retain their previous totals. This lets the receiver refresh - // only the resources whose available capacity changed. - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (present[i] && - (proposed[i] < e.grants[i] || proposed[i] < e.consumed[i])) + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (present[i] && (proposed[i] < entry.grants[i] || + proposed[i] < entry.consumed[i])) { return Status::InvalidArgument( "decreasing or under-consumed grant" LOC_MARK); - bool gap = e.has_update && u.sequence > e.last_sequence + 1; - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (present[i]) e.grants[i] = proposed[i]; - e.last_sequence = u.sequence; - e.has_update = true; + } + } + + const bool gap = + entry.has_update && update.sequence > entry.last_sequence + 1; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (present[i]) entry.grants[i] = proposed[i]; + } + entry.last_sequence = update.sequence; + entry.has_update = true; + const uint64_t ttl_ns = + static_cast(update.freshness_ttl_ms) * 1000000ULL; + const uint64_t now = nowNanos(); + entry.expires_at_ns = ttl_ns > std::numeric_limits::max() - now + ? std::numeric_limits::max() + : now + ttl_ns; disposition = gap ? CreditUpdateDisposition::SequenceGap : CreditUpdateDisposition::Applied; return Status::OK(); } -Status SenderCreditLedger::tryReserve(const CreditKey& k, - const CreditCharge& c) { - std::array n; - CHECK_STATUS(normalize(c, n)); +Status SenderCreditLedger::tryReserve(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || !it->second.has_update) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit unavailable" LOC_MARK); - auto& e = it->second; - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (e.consumed[i] > e.grants[i] || n[i] > e.grants[i] - e.consumed[i]) + } + CHECK_STATUS(requireFresh(it->second)); + auto& entry = it->second; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (entry.consumed[i] > entry.grants[i] || + normalized[i] > entry.grants[i] - entry.consumed[i]) { return Status::TooManyRequests("insufficient credit" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) e.consumed[i] += n[i]; + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + entry.consumed[i] += normalized[i]; + } return Status::OK(); } -Status SenderCreditLedger::rollbackReservation(const CreditKey& k, - const CreditCharge& c) { - std::array n; - CHECK_STATUS(normalize(c, n)); +Status SenderCreditLedger::rollbackReservation(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end()) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit session inactive" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) - if (n[i] > it->second.consumed[i]) + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (it->second.released[i] > it->second.consumed[i] || + normalized[i] > it->second.consumed[i] - it->second.released[i]) { return Status::InvalidArgument( "credit rollback underflow" LOC_MARK); - for (size_t i = 0; i < kCreditResourceCount; ++i) - it->second.consumed[i] -= n[i]; + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (normalized[i] > + std::numeric_limits::max() - it->second.released[i]) { + return Status::InvalidArgument("credit rollback overflow" LOC_MARK); + } + it->second.released[i] += normalized[i]; + } + return Status::OK(); +} + +Status SenderCreditLedger::releaseCommitted(const CreditKey& key, + const CreditCharge& charge) { + std::array normalized; + CHECK_STATUS(normalize(charge, normalized)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (it->second.released[i] > it->second.consumed[i] || + normalized[i] > it->second.consumed[i] - it->second.released[i]) { + return Status::InvalidArgument( + "credit release exceeds committed consumption" LOC_MARK); + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (normalized[i] > + std::numeric_limits::max() - it->second.released[i]) { + return Status::InvalidArgument("credit release overflow" LOC_MARK); + } + it->second.released[i] += normalized[i]; + } + return Status::OK(); +} + +Status SenderCreditLedger::prepareExchange( + const CreditKey& key, const CreditCharge& next_charge, + ReceiverCreditExchangeRequestV1& request) { + std::array normalized; + CHECK_STATUS(normalize(next_charge, normalized)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + auto& entry = it->second; + if (entry.report_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "credit report sequence exhausted" LOC_MARK); + } + + request = ReceiverCreditExchangeRequestV1{}; + request.expected_receiver_session_id = key.receiver_session; + request.epoch = entry.epoch; + request.sender_instance_id = key.sender_instance; + request.qos_class = key.qos_class; + request.known_update_sequence = entry.last_sequence; + request.report_sequence = ++entry.report_sequence; + request.released_totals.reserve(kCreditResourceCount); + request.desired_grant_totals.reserve(next_charge.resources.size()); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const auto resource = static_cast(i + 1); + request.released_totals.push_back({resource, entry.released[i]}); + if (normalized[i] == 0) continue; + if (normalized[i] > + std::numeric_limits::max() - entry.consumed[i]) { + return Status::InvalidArgument( + "desired credit total overflow" LOC_MARK); + } + request.desired_grant_totals.push_back( + {resource, entry.consumed[i] + normalized[i]}); + } + return Status::OK(); +} + +Status SenderCreditLedger::prepareReleaseExchange( + const CreditKey& key, ReceiverCreditExchangeRequestV1& request) { + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + auto& entry = it->second; + if (entry.report_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "credit report sequence exhausted" LOC_MARK); + } + + request = ReceiverCreditExchangeRequestV1{}; + request.expected_receiver_session_id = key.receiver_session; + request.epoch = entry.epoch; + request.sender_instance_id = key.sender_instance; + request.qos_class = key.qos_class; + request.known_update_sequence = entry.last_sequence; + request.report_sequence = ++entry.report_sequence; + request.released_totals.reserve(kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + const auto resource = static_cast(i + 1); + request.released_totals.push_back({resource, entry.released[i]}); + if (entry.consumed[i] != 0) { + request.desired_grant_totals.push_back( + {resource, entry.consumed[i]}); + } + } + if (request.desired_grant_totals.empty()) { + return Status::InvalidEntry( + "credit release exchange has no consumption" LOC_MARK); + } return Status::OK(); } -Status SenderCreditLedger::available(const CreditKey& k, CreditResource r, - uint64_t& v) const { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); +Status SenderCreditLedger::available(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end() || !it->second.has_update) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit unavailable" LOC_MARK); - v = it->second.grants[i] - it->second.consumed[i]; + } + CHECK_STATUS(requireFresh(it->second)); + value = it->second.grants[index] - it->second.consumed[index]; + return Status::OK(); +} + +Status SenderCreditLedger::consumed(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + return Status::InvalidEntry("credit session inactive" LOC_MARK); + } + value = it->second.consumed[index]; return Status::OK(); } -Status SenderCreditLedger::consumed(const CreditKey& k, CreditResource r, - uint64_t& v) const { - size_t i = 0; - CHECK_STATUS(resourceIndex(r, i)); +Status SenderCreditLedger::released(const CreditKey& key, + CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(resourceIndex(resource, index)); std::lock_guard lock(mutex_); - auto it = entries_.find(k); - if (it == entries_.end()) + auto it = entries_.find(key); + if (it == entries_.end()) { return Status::InvalidEntry("credit session inactive" LOC_MARK); - v = it->second.consumed[i]; + } + value = it->second.released[index]; + return Status::OK(); +} + +ReceiverCreditAuthority::ReceiverCreditAuthority(ReceiverSessionId session, + uint64_t epoch, + ReceiverCreditLimits limits) + : session_(session), epoch_(epoch), limits_(std::move(limits)) { + if (session_.empty() || epoch_ == 0) { + status_ = Status::InvalidArgument( + "invalid receiver credit authority identity" LOC_MARK); + } else { + status_ = validateLimits(limits_); + } +} + +Status ReceiverCreditAuthority::validateLimits( + const ReceiverCreditLimits& limits) { + if (limits.freshness_ttl_ms == 0 || limits.max_senders == 0 || + limits.max_entries < limits.max_senders) { + return Status::InvalidArgument( + "invalid receiver credit limits" LOC_MARK); + } + bool any_capacity = false; + for (const auto capacity : limits.capacity) any_capacity |= capacity != 0; + if (!any_capacity) { + return Status::InvalidArgument( + "receiver credit capacity is empty" LOC_MARK); + } + return Status::OK(); +} + +ReceiverCreditAdvertV1 ReceiverCreditAuthority::advert() const { + ReceiverCreditAdvertV1 result; + result.receiver_session_id = session_; + result.epoch = epoch_; + result.freshness_ttl_ms = limits_.freshness_ttl_ms; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (limits_.capacity[i] != 0) { + const auto resource = static_cast(i + 1); + result.resources.push_back(resource); + result.capacities.push_back({resource, limits_.capacity[i]}); + } + } + return result; +} + +Status ReceiverCreditAuthority::normalizeCounters( + const std::vector& counters, + std::array& normalized, + std::array& present) { + normalized.fill(0); + present.fill(false); + if (counters.size() > kCreditResourceCount) { + return Status::InvalidArgument("too many credit counters" LOC_MARK); + } + for (const auto& counter : counters) { + size_t index = 0; + CHECK_STATUS(indexFor(counter.resource, index)); + if (present[index]) { + return Status::InvalidArgument( + "duplicate credit counter resource" LOC_MARK); + } + present[index] = true; + normalized[index] = counter.total; + } + return Status::OK(); +} + +void ReceiverCreditAuthority::fillReply( + const CreditKey& key, const Entry& entry, + ReceiverCreditExchangeReplyV1& reply) const { + reply = ReceiverCreditExchangeReplyV1{}; + reply.sender_instance_id = key.sender_instance; + reply.update.qos_class = key.qos_class; + reply.update.receiver_session_id = session_; + reply.update.epoch = epoch_; + reply.update.sequence = entry.update_sequence; + reply.update.freshness_ttl_ms = limits_.freshness_ttl_ms; + reply.update.grants.reserve(kCreditResourceCount); + for (size_t i = 0; i < kCreditResourceCount; ++i) { + reply.update.grants.push_back( + {static_cast(i + 1), entry.grants[i]}); + } +} + +Status ReceiverCreditAuthority::exchange( + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + CHECK_STATUS(status_); + if (request.schema_version != kReceiverCreditProtocolVersion || + request.flags != 0 || request.epoch != epoch_ || + !(request.expected_receiver_session_id == session_) || + request.sender_instance_id.empty() || request.qos_class > 2 || + request.report_sequence == 0 || request.desired_grant_totals.empty()) { + return Status::InvalidArgument( + "invalid receiver credit exchange identity/header" LOC_MARK); + } + + std::array released{}, desired{}; + std::array released_present{}, + desired_present{}; + CHECK_STATUS( + normalizeCounters(request.released_totals, released, released_present)); + CHECK_STATUS(normalizeCounters(request.desired_grant_totals, desired, + desired_present)); + + CreditKey key{request.expected_receiver_session_id, + request.sender_instance_id, request.qos_class}; + std::lock_guard lock(mutex_); + auto it = entries_.find(key); + if (it == entries_.end()) { + if (request.known_update_sequence != 0) { + return Status::InvalidArgument( + "new sender knows an impossible credit update" LOC_MARK); + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (released_present[i] && released[i] != 0) { + return Status::InvalidArgument( + "new sender reports impossible credit release" LOC_MARK); + } + } + if (entries_.size() >= limits_.max_entries) { + return Status::TooManyRequests( + "receiver credit replay-history limit" LOC_MARK); + } + it = entries_.emplace(key, Entry{}).first; + } + auto& entry = it->second; + if (request.known_update_sequence > entry.update_sequence) { + return Status::InvalidArgument( + "sender knows an impossible credit update" LOC_MARK); + } + if (request.report_sequence <= entry.last_report_sequence) { + fillReply(key, entry, reply); + return Status::OK(); + } + if (entry.update_sequence == std::numeric_limits::max()) { + return Status::InternalError( + "receiver credit update sequence exhausted" LOC_MARK); + } + + // Validate the entire cumulative release before mutating any resource. + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + if (released[i] < entry.released[i] || released[i] > entry.grants[i]) { + return Status::InvalidArgument( + "invalid cumulative credit release" LOC_MARK); + } + } + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + const uint64_t delta = released[i] - entry.released[i]; + if (delta > outstanding_[i]) { + return Status::InternalError( + "receiver credit accounting underflow" LOC_MARK); + } + } + + // Apply releases first, then atomically decide whether every requested + // resource can reach its absolute desired total. + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!released_present[i]) continue; + outstanding_[i] -= released[i] - entry.released[i]; + entry.released[i] = released[i]; + } + + std::array grant_delta{}; + bool can_grant_all = true; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + if (!desired_present[i] || desired[i] <= entry.grants[i]) continue; + grant_delta[i] = desired[i] - entry.grants[i]; + const uint64_t free = limits_.capacity[i] - outstanding_[i]; + if (grant_delta[i] > free) can_grant_all = false; + } + const bool requests_new_grant = + std::any_of(grant_delta.begin(), grant_delta.end(), + [](uint64_t delta) { return delta != 0; }); + if (can_grant_all && requests_new_grant) { + std::unordered_set + active_senders; + active_senders.reserve(std::min(entries_.size(), limits_.max_senders)); + for (const auto& [entry_key, tracked] : entries_) { + bool has_outstanding = false; + for (size_t i = 0; i < kCreditResourceCount; ++i) { + has_outstanding |= tracked.grants[i] > tracked.released[i]; + } + if (has_outstanding) { + active_senders.insert(entry_key.sender_instance); + } + } + if (!active_senders.count(request.sender_instance_id) && + active_senders.size() >= limits_.max_senders) { + can_grant_all = false; + } + } + if (can_grant_all) { + for (size_t i = 0; i < kCreditResourceCount; ++i) { + entry.grants[i] += grant_delta[i]; + outstanding_[i] += grant_delta[i]; + } + } + + entry.last_report_sequence = request.report_sequence; + ++entry.update_sequence; + fillReply(key, entry, reply); + return Status::OK(); +} + +Status ReceiverCreditAuthority::outstanding(CreditResource resource, + uint64_t& value) const { + size_t index = 0; + CHECK_STATUS(indexFor(resource, index)); + std::lock_guard lock(mutex_); + value = outstanding_[index]; return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 305844d94f..4528e527d7 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -57,6 +57,33 @@ constexpr uint8_t kRedisDefaultDbIndex = 0; // the counter), so only a permanently failing poll or queue retire gets here. constexpr size_t kMaxReclaimAttempts = 4096; constexpr int kMaxHpTcpMetadataRefreshRetries = 1; +uint64_t randomIdentityWord() { + try { + static thread_local std::random_device random; + return (static_cast(random()) << 32) ^ random() ^ + getCurrentTimeInNano(); + } catch (const std::exception&) { + // random_device is allowed to be unavailable on unusual libstdc++ + // targets. Preserve startup while still mixing a process-local nonce + // into the clock-seeded fallback generator. + static std::atomic nonce{1}; + auto& random = SimpleRandom::Get(); + return (static_cast(random.next()) << 32) ^ random.next() ^ + nonce.fetch_add(1, std::memory_order_relaxed); + } +} + +SenderInstanceId makeSenderInstanceId() { + SenderInstanceId id{randomIdentityWord(), randomIdentityWord()}; + if (id.empty()) id.low = 1; + return id; +} + +ReceiverSessionId makeReceiverSessionId() { + ReceiverSessionId id{randomIdentityWord(), randomIdentityWord()}; + if (id.empty()) id.low = 1; + return id; +} } // namespace struct Batch { @@ -360,6 +387,11 @@ Status TransferEngineImpl::setupLocalSegment() { segment.rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); auto& detail = std::get(segment.detail); detail.topology = *(topology_.get()); + if (receiver_credit_authority_) { + detail.receiver_credit = receiver_credit_authority_->advert(); + } else { + detail.receiver_credit.reset(); + } return Status::OK(); })); local_segment_tracker_ = std::make_unique(manager); @@ -395,6 +427,15 @@ Status TransferEngineImpl::construct() { notify_rpc_fallback_ = conf_->get("notification/rpc_fallback", true); enable_progress_worker_ = conf_->get("enable_progress_worker", false); runtime_queue_config_.enabled = conf_->get("enable_runtime_queue", false); + const bool receiver_credit_legacy_enable = + conf_->get("receiver_credit/enable", false); + receiver_credit_config_.sender_enabled = conf_->get( + "receiver_credit/sender_enable", receiver_credit_legacy_enable); + receiver_credit_config_.receiver_enabled = conf_->get( + "receiver_credit/receiver_enable", receiver_credit_legacy_enable); + if (receiver_credit_config_.sender_enabled) { + runtime_queue_config_.enabled = true; + } if (runtime_queue_config_.enabled) enable_progress_worker_ = true; runtime_queue_config_.limits.max_outstanding_owners = conf_->get("runtime_queue/max_outstanding_owners", 1024UL); @@ -417,6 +458,59 @@ Status TransferEngineImpl::construct() { runtime_queue_config_.progress_fallback_interval = std::chrono::microseconds( conf_->get("runtime_queue/progress_fallback_interval_us", 50000UL)); + + receiver_credit_config_.limits.capacity[0] = + conf_->get("receiver_credit/data_bytes", + static_cast( + runtime_queue_config_.limits.max_outstanding_bytes)); + receiver_credit_config_.limits.capacity[1] = + conf_->get("receiver_credit/request_slots", + static_cast( + runtime_queue_config_.limits.max_outstanding_owners)); + receiver_credit_config_.limits.capacity[2] = 0; + receiver_credit_config_.limits.capacity[3] = 0; + receiver_credit_config_.limits.freshness_ttl_ms = + conf_->get("receiver_credit/freshness_ttl_ms", uint32_t{1000}); + receiver_credit_config_.limits.max_senders = + conf_->get("receiver_credit/max_senders", size_t{1024}); + receiver_credit_config_.limits.max_entries = + conf_->get("receiver_credit/max_authority_entries", size_t{4096}); + receiver_credit_config_.max_sender_entries = + conf_->get("receiver_credit/max_sender_entries", size_t{4096}); + receiver_credit_config_.max_freshness_ttl_ms = + conf_->get("receiver_credit/max_freshness_ttl_ms", uint32_t{60000}); + if (conf_->get("receiver_credit/staging_slots", uint64_t{0}) != 0 || + conf_->get("receiver_credit/consumer_slots", uint64_t{0}) != 0) { + return Status::InvalidArgument( + "staging/consumer receiver-credit resources are not supported " + "by the PR6 runtime gate" LOC_MARK); + } + if (receiver_credit_config_.sender_enabled) { + if (receiver_credit_config_.max_sender_entries == 0 || + receiver_credit_config_.max_freshness_ttl_ms == 0 || + runtime_queue_config_.progress_fallback_interval.count() <= 0) { + return Status::InvalidArgument( + "receiver-credit sender requires non-zero entries, TTL " + "limit, and progress fallback interval" LOC_MARK); + } + } + if (receiver_credit_config_.receiver_enabled) { + if (receiver_credit_config_.limits.capacity[0] == 0 || + receiver_credit_config_.limits.capacity[1] == 0 || + receiver_credit_config_.limits.max_senders == 0 || + receiver_credit_config_.limits.max_entries < + receiver_credit_config_.limits.max_senders) { + return Status::InvalidArgument( + "receiver credit requires non-zero bytes, request slots, " + "and a replay-history limit no smaller than the sender " + "limit" LOC_MARK); + } + if (receiver_credit_config_.limits.freshness_ttl_ms > + receiver_credit_config_.max_freshness_ttl_ms) { + return Status::InvalidArgument( + "receiver credit TTL exceeds sender safety limit" LOC_MARK); + } + } if (runtime_queue_config_.enabled && (runtime_queue_config_.max_dispatch_owners == 0 || runtime_queue_config_.max_dispatch_bytes == 0)) { @@ -425,6 +519,15 @@ Status TransferEngineImpl::construct() { } runtime_queue_ = std::make_unique( runtime_queue_config_.limits); + sender_instance_id_ = makeSenderInstanceId(); + sender_credit_ledger_ = std::make_unique( + std::max(receiver_credit_config_.max_sender_entries, 1), + std::max(receiver_credit_config_.max_freshness_ttl_ms, 1)); + if (receiver_credit_config_.receiver_enabled) { + receiver_credit_authority_ = std::make_unique( + makeReceiverSessionId(), 1, receiver_credit_config_.limits); + CHECK_STATUS(receiver_credit_authority_->status()); + } if (!hostname_.empty()) CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_)); else @@ -437,6 +540,14 @@ Status TransferEngineImpl::construct() { metadata_ = std::make_shared(metadata_type, metadata_servers, this); + if (receiver_credit_authority_) { + metadata_->setReceiverCreditCallback( + [this](const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + return receiver_credit_authority_->exchange(request, reply); + }); + } + CHECK_STATUS(metadata_->start(port_, ipv6_, rpc_server_threads)); if (metadata_type == "p2p") @@ -610,6 +721,28 @@ Status TransferEngineImpl::deconstruct() { << status.ToString(); } } + // Only a reservation that never reached a transport can be returned here. + // A remaining committed owner is not terminal, and the UB DMA fence is + // established later while its transport is destroyed. Conservatively + // abandon that authorization in the old receiver session instead of + // making it available while DMA may still be active. + size_t abandoned_receiver_credits = 0; + for (auto& [_, queued] : queued_owners_) { + if (!queued.credit.active) continue; + if (!queued.credit.committed) { + rollbackReceiverCredit(queued.credit); + } else { + queued.credit = ReceiverCreditReservation{}; + ++abandoned_receiver_credits; + } + } + if (abandoned_receiver_credits != 0) { + LOG(WARNING) << "Conservatively abandoned " + << abandoned_receiver_credits + << " in-flight receiver-credit reservations during " + "shutdown"; + } + progressReceiverCreditReleases(); // Destroy staging_proxy_ first: its destructor calls back into // unregisterLocalMemory/freeLocalMemory, which require @@ -1698,6 +1831,10 @@ struct RequestBoundaryInfo { std::optional source_key; std::optional target_key; uint64_t max_merge_bytes{std::numeric_limits::max()}; + // A required receiver-credit advert limits the size of one physical + // owner. Request merging must not turn individually admissible requests + // into an owner that can never acquire credit. + std::optional receiver_credit_data_capacity; }; struct MergeResult { @@ -1720,6 +1857,8 @@ struct TransferEngineImpl::PreparedSubmit { std::vector derived_task_ids; Request request{}; SelectionResult route{}; + bool receiver_credit_required{false}; + bool receiver_credit_exempt{false}; bool staging{false}; std::vector staging_params; }; @@ -1820,6 +1959,15 @@ MergeResult mergeRequests(const std::vector& requests, if (last.req.length + curr.req.length > last.boundary.max_merge_bytes) { return false; } + const size_t merged_length = last.req.length + curr.req.length; + if (last.boundary.receiver_credit_data_capacity != + curr.boundary.receiver_credit_data_capacity) { + return false; + } + if (last.boundary.receiver_credit_data_capacity && + merged_length > *last.boundary.receiver_credit_data_capacity) { + return false; + } uint64_t last_source_end = 0; uint64_t last_target_end = 0; @@ -1898,6 +2046,21 @@ std::vector resolveRequestBoundaries( std::min(hp_tcp_limit, endpoint.max_transfer_bytes); } } + std::optional receiver_credit_data_capacity; + if (target_desc->type == SegmentType::Memory && + target_desc->getMemory().receiver_credit) { + // Zero deliberately disables merging for malformed + // adverts; prepareSubmit validates and rejects them with + // the protocol-specific error below. + receiver_credit_data_capacity = 0; + for (const auto& capacity : + target_desc->getMemory().receiver_credit->capacities) { + if (capacity.resource == CreditResource::DataBytes) { + receiver_credit_data_capacity = capacity.total; + break; + } + } + } bool any_missing = false; for (size_t i : idxs) { const auto& r = requests[i]; @@ -1907,6 +2070,8 @@ std::vector resolveRequestBoundaries( r.transport_hint == UNSPEC)) { boundaries[i].max_merge_bytes = hp_tcp_limit; } + boundaries[i].receiver_credit_data_capacity = + receiver_credit_data_capacity; auto* buffer = target_desc->findBuffer(r.target_offset, r.length); if (!buffer) { @@ -2090,7 +2255,8 @@ SelectionResult TransferEngineImpl::resolveTransport(const Request& req, Status TransferEngineImpl::prepareSubmit( Batch* batch, const std::vector& request_list, - PreparedSubmit& prepared, bool require_post_submit_cancellation) { + QueueOwnerKind owner_kind, PreparedSubmit& prepared, + bool require_post_submit_cancellation) { if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); for (size_t i = 0; i < request_list.size(); ++i) { auto st = validateTransportHint(request_list[i], i); @@ -2171,6 +2337,49 @@ Status TransferEngineImpl::prepareSubmit( "transport" LOC_MARK); } } + owner.receiver_credit_exempt = + owner_kind == QueueOwnerKind::StagingInternal; + if (owner_kind != QueueOwnerKind::StagingInternal) { + ReceiverCreditTarget credit_target; + CHECK_STATUS( + resolveReceiverCreditTarget(owner.request, credit_target)); + owner.receiver_credit_required = credit_target.required; + // Scatter bypasses the runtime queue to preserve synchronous + // cancellation. It cannot bypass a peer's required credit gate. + if (credit_target.required && require_post_submit_cancellation) { + return Status::NotImplemented( + "scatter transfer does not support required receiver " + "credit" LOC_MARK); + } + if (credit_target.required && + !receiver_credit_config_.sender_enabled) { + return Status::InvalidArgument( + "peer requires receiver credit but the sender gate is " + "disabled" LOC_MARK); + } + CHECK_STATUS( + validateReceiverCreditCapacity(owner.request, credit_target)); + } else { + const bool has_credit_fence = + owner.request.receiver_credit_session_high != 0 || + owner.request.receiver_credit_session_low != 0 || + owner.request.receiver_credit_epoch != 0; + if (has_credit_fence) { + ReceiverCreditTarget credit_target; + CHECK_STATUS( + resolveReceiverCreditTarget(owner.request, credit_target)); + if (!credit_target.required || + credit_target.advert.receiver_session_id.high != + owner.request.receiver_credit_session_high || + credit_target.advert.receiver_session_id.low != + owner.request.receiver_credit_session_low || + credit_target.advert.epoch != + owner.request.receiver_credit_epoch) { + return Status::InvalidEntry( + "staging receiver-credit session changed" LOC_MARK); + } + } + } owner.staging_params = owner.route.staging_params; if (owner.staging_params.empty() && (owner.route.transport == TCP || owner.route.transport == HP_TCP)) { @@ -2214,6 +2423,354 @@ void TransferEngineImpl::attachProgressNotifier( }; } +Status TransferEngineImpl::resolveReceiverCreditTarget( + const Request& request, ReceiverCreditTarget& target) const { + target = ReceiverCreditTarget{}; + SegmentDescRef segment; + if (request.target_id == LOCAL_SEGMENT_ID) { + segment = metadata_->segmentManager().getLocal(); + } else { + CHECK_STATUS(metadata_->segmentManager().getRemoteCached( + segment, request.target_id)); + } + if (!segment || segment->type != SegmentType::Memory) return Status::OK(); + const auto& memory = segment->getMemory(); + if (!memory.receiver_credit) return Status::OK(); + + const auto& advert = *memory.receiver_credit; + if (advert.schema_version != kReceiverCreditProtocolVersion || + advert.flags != kReceiverCreditRequired || + advert.receiver_session_id.empty() || advert.epoch == 0 || + advert.freshness_ttl_ms == 0 || + advert.freshness_ttl_ms > + receiver_credit_config_.max_freshness_ttl_ms || + advert.resources.empty() || + advert.resources.size() > kCreditResourceCount || + advert.capacities.size() != advert.resources.size()) { + return Status::InvalidArgument( + "invalid receiver-credit advertisement" LOC_MARK); + } + + std::array resources{}; + for (const auto resource : advert.resources) { + const auto raw = static_cast(resource); + if (raw < 1 || raw > kCreditResourceCount || resources[raw - 1]) { + return Status::InvalidArgument( + "invalid receiver-credit resource advertisement" LOC_MARK); + } + resources[raw - 1] = true; + } + std::array capacities{}; + for (const auto& capacity : advert.capacities) { + const auto raw = static_cast(capacity.resource); + if (raw < 1 || raw > kCreditResourceCount || capacity.total == 0 || + capacities[raw - 1] || !resources[raw - 1]) { + return Status::InvalidArgument( + "invalid receiver-credit capacity advertisement" LOC_MARK); + } + capacities[raw - 1] = true; + } + if (capacities != resources) { + return Status::InvalidArgument( + "receiver-credit resources/capacities mismatch" LOC_MARK); + } + if (!resources[static_cast(CreditResource::DataBytes) - 1] || + !resources[static_cast(CreditResource::RequestSlots) - 1]) { + return Status::InvalidArgument( + "receiver credit lacks bytes/request resources" LOC_MARK); + } + if (request.target_id != LOCAL_SEGMENT_ID && + segment->rpc_server_addr.empty()) { + return Status::InvalidArgument( + "receiver credit peer has no RPC address" LOC_MARK); + } + target.required = true; + target.advert = advert; + target.rpc_server_addr = segment->rpc_server_addr; + return Status::OK(); +} + +Status TransferEngineImpl::validateReceiverCreditCapacity( + const Request& request, const ReceiverCreditTarget& target) const { + if (!target.required) return Status::OK(); + uint64_t data_bytes = 0; + uint64_t request_slots = 0; + for (const auto& capacity : target.advert.capacities) { + if (capacity.resource == CreditResource::DataBytes) { + data_bytes = capacity.total; + } else if (capacity.resource == CreditResource::RequestSlots) { + request_slots = capacity.total; + } + } + if (request_slots == 0 || request.length > data_bytes) { + return Status::TooManyRequests( + "request exceeds receiver-credit capacity" LOC_MARK); + } + return Status::OK(); +} + +Status TransferEngineImpl::exchangeReceiverCredit( + SegmentID target_id, const std::string& rpc_server_addr, + const ReceiverCreditExchangeRequestV1& request, + ReceiverCreditExchangeReplyV1& reply) { + if (target_id == LOCAL_SEGMENT_ID) { + if (!receiver_credit_authority_) { + return Status::InvalidEntry( + "local receiver-credit authority is unavailable" LOC_MARK); + } + return receiver_credit_authority_->exchange(request, reply); + } + return ControlClient::exchangeReceiverCredit(rpc_server_addr, request, + reply); +} + +Status TransferEngineImpl::acquireReceiverCredit( + const Request& request, ReceiverCreditReservation& reservation, + bool& acquired) { + acquired = false; + if (reservation.active) { + return Status::InternalError( + "receiver credit attempt already has a reservation" LOC_MARK); + } + + ReceiverCreditTarget target; + CHECK_STATUS(resolveReceiverCreditTarget(request, target)); + if (!target.required) { + acquired = true; + return Status::OK(); + } + CHECK_STATUS(validateReceiverCreditCapacity(request, target)); + if (!receiver_credit_config_.sender_enabled || !sender_credit_ledger_) { + return Status::InvalidEntry( + "receiver-credit sender gate is disabled" LOC_MARK); + } + + const uint32_t qos_class = static_cast( + std::clamp(request.priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW))); + CreditKey key{target.advert.receiver_session_id, sender_instance_id_, + qos_class}; + CreditCharge charge{{{CreditResource::DataBytes, request.length}, + {CreditResource::RequestSlots, 1}}}; + observeReceiverCreditSession(request.target_id, + target.advert.receiver_session_id); + CHECK_STATUS(sender_credit_ledger_->activate(key, target.advert.epoch)); + + auto reserve_status = sender_credit_ledger_->tryReserve(key, charge); + if (!reserve_status.ok()) { + if (!reserve_status.IsTooManyRequests() && + !reserve_status.IsInvalidEntry()) { + return reserve_status; + } + + ReceiverCreditExchangeRequestV1 exchange_request; + CHECK_STATUS(sender_credit_ledger_->prepareExchange(key, charge, + exchange_request)); + ReceiverCreditExchangeReplyV1 exchange_reply; + auto exchange_status = + exchangeReceiverCredit(request.target_id, target.rpc_server_addr, + exchange_request, exchange_reply); + if (!exchange_status.ok()) { + if (request.target_id != LOCAL_SEGMENT_ID) { + (void)metadata_->segmentManager().invalidateRemote( + request.target_id); + } + LOG(WARNING) << "Receiver-credit exchange deferred: " + << exchange_status.ToString(); + return Status::OK(); + } + if (!(exchange_reply.sender_instance_id == sender_instance_id_)) { + return Status::InvalidArgument( + "receiver-credit reply sender mismatch" LOC_MARK); + } + if (exchange_reply.flags != 0) { + return Status::InvalidArgument( + "receiver-credit demand reply flags" LOC_MARK); + } + CreditUpdateDisposition disposition; + CHECK_STATUS(sender_credit_ledger_->applyUpdate( + key, exchange_reply.update, disposition)); + pending_credit_releases_.erase(key); + reserve_status = sender_credit_ledger_->tryReserve(key, charge); + if (reserve_status.IsTooManyRequests() || + reserve_status.IsInvalidEntry()) { + return Status::OK(); + } + CHECK_STATUS(reserve_status); + } + + reservation.active = true; + reservation.committed = false; + reservation.key = key; + reservation.charge = std::move(charge); + reservation.target_id = request.target_id; + reservation.rpc_server_addr = std::move(target.rpc_server_addr); + reservation.epoch = target.advert.epoch; + acquired = true; + return Status::OK(); +} + +void TransferEngineImpl::rollbackReceiverCredit( + ReceiverCreditReservation& reservation) { + if (!reservation.active) return; + if (!reservation.committed) { + const auto key = reservation.key; + const auto target_id = reservation.target_id; + const auto rpc_server_addr = reservation.rpc_server_addr; + const auto epoch = reservation.epoch; + auto status = sender_credit_ledger_->rollbackReservation( + reservation.key, reservation.charge); + reservation = ReceiverCreditReservation{}; + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit rollback failed: " + << status.ToString(); + return; + } + scheduleReceiverCreditRelease(key, target_id, rpc_server_addr, epoch); + (void)tryReportReceiverCreditRelease(key); + return; + } + reservation = ReceiverCreditReservation{}; +} + +void TransferEngineImpl::releaseReceiverCredit( + ReceiverCreditReservation& reservation) { + if (!reservation.active) return; + if (!reservation.committed) { + rollbackReceiverCredit(reservation); + return; + } + + const auto key = reservation.key; + const auto target_id = reservation.target_id; + const auto rpc_server_addr = reservation.rpc_server_addr; + const auto epoch = reservation.epoch; + auto status = sender_credit_ledger_->releaseCommitted(reservation.key, + reservation.charge); + reservation = ReceiverCreditReservation{}; + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit terminal release failed: " + << status.ToString(); + return; + } + + scheduleReceiverCreditRelease(key, target_id, rpc_server_addr, epoch); + (void)tryReportReceiverCreditRelease(key); +} + +void TransferEngineImpl::scheduleReceiverCreditRelease( + const CreditKey& key, SegmentID target_id, + const std::string& rpc_server_addr, uint64_t epoch) { + auto& pending = pending_credit_releases_[key]; + pending.target_id = target_id; + pending.rpc_server_addr = rpc_server_addr; + pending.epoch = epoch; + pending.retry_at = std::chrono::steady_clock::now(); +} + +bool TransferEngineImpl::tryReportReceiverCreditRelease(const CreditKey& key) { + auto pending_it = pending_credit_releases_.find(key); + if (pending_it == pending_credit_releases_.end()) return true; + auto& pending = pending_it->second; + if (std::chrono::steady_clock::now() < pending.retry_at) return false; + + auto defer_retry = [&](const Status& status) { + ++pending.attempts; + const uint32_t shift = std::min(pending.attempts - 1, 4); + const auto delay = std::chrono::milliseconds(50U << shift); + pending.retry_at = std::chrono::steady_clock::now() + delay; + LOG(WARNING) << "Receiver-credit release exchange deferred (attempt " + << pending.attempts << "): " << status.ToString(); + return false; + }; + + ReceiverCreditExchangeRequestV1 request; + auto status = sender_credit_ledger_->prepareReleaseExchange(key, request); + if (!status.ok()) { + return defer_retry(status); + } + ReceiverCreditExchangeReplyV1 reply; + status = exchangeReceiverCredit(pending.target_id, pending.rpc_server_addr, + request, reply); + if (!status.ok()) { + return defer_retry(status); + } + if (!(reply.sender_instance_id == sender_instance_id_)) { + return defer_retry(Status::InvalidArgument( + "receiver-credit release reply sender mismatch" LOC_MARK)); + } + if (reply.flags != 0) { + return defer_retry(Status::InvalidArgument( + "receiver-credit release reply flags" LOC_MARK)); + } + CreditUpdateDisposition disposition; + status = sender_credit_ledger_->applyUpdate(key, reply.update, disposition); + if (!status.ok()) { + return defer_retry(status); + } + pending_credit_releases_.erase(pending_it); + notifyRuntimeQueueReady(); + return true; +} + +void TransferEngineImpl::progressReceiverCreditReleases() { + const auto now = std::chrono::steady_clock::now(); + std::vector due; + due.reserve(pending_credit_releases_.size()); + for (const auto& [key, pending] : pending_credit_releases_) { + if (pending.retry_at <= now) due.push_back(key); + } + for (const auto& key : due) (void)tryReportReceiverCreditRelease(key); +} + +void TransferEngineImpl::observeReceiverCreditSession( + SegmentID target_id, const ReceiverSessionId& session) { + auto [it, inserted] = + receiver_credit_sessions_.try_emplace(target_id, session); + if (inserted || it->second == session) return; + + const auto stale_session = it->second; + auto status = sender_credit_ledger_->deactivateSession(stale_session); + if (!status.ok()) { + LOG(ERROR) << "Receiver-credit stale session cleanup failed: " + << status.ToString(); + } + for (auto pending = pending_credit_releases_.begin(); + pending != pending_credit_releases_.end();) { + if (pending->first.receiver_session == stale_session) { + pending = pending_credit_releases_.erase(pending); + } else { + ++pending; + } + } + it->second = session; +} + +void TransferEngineImpl::bindReceiverCreditFence( + Transport::SubBatchRef sub_batch, + const ReceiverCreditReservation& reservation) const { + if (!sub_batch) return; + sub_batch->receiver_credit_session_high = 0; + sub_batch->receiver_credit_session_low = 0; + sub_batch->receiver_credit_epoch = 0; + if (!reservation.active) return; + sub_batch->receiver_credit_session_high = + reservation.key.receiver_session.high; + sub_batch->receiver_credit_session_low = + reservation.key.receiver_session.low; + sub_batch->receiver_credit_epoch = reservation.epoch; +} + +void TransferEngineImpl::bindReceiverCreditFence( + Transport::SubBatchRef sub_batch, const Request& request) const { + if (!sub_batch) return; + sub_batch->receiver_credit_session_high = + request.receiver_credit_session_high; + sub_batch->receiver_credit_session_low = + request.receiver_credit_session_low; + sub_batch->receiver_credit_epoch = request.receiver_credit_epoch; +} + Status TransferEngineImpl::commitPreparedSubmit( Batch* batch, const PreparedSubmit& prepared) { if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); @@ -2487,6 +3044,8 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, queued.batch = batch; queued.owner_task_id = prepared.owners[i].owner_task_id; queued.byte_charge = prepared.owners[i].request.length; + queued.receiver_credit_exempt = + prepared.owners[i].receiver_credit_exempt; queued.public_task_ids.push_back(prepared.owners[i].owner_task_id); queued.public_task_ids.insert( queued.public_task_ids.end(), @@ -2511,6 +3070,7 @@ Status TransferEngineImpl::finishQueuedOwner( "runtime dispatch window accounting underflow" LOC_MARK); } } + releaseReceiverCredit(queued.credit); CHECK_STATUS(runtime_queue_->complete(owner_id, terminal_status)); if (queued.in_dispatch_window) { --dispatch_inflight_owners_; @@ -2611,7 +3171,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (queued_it == queued_owners_.end()) { return Status::InternalError("queued owner metadata missing" LOC_MARK); } - const auto queued = queued_it->second; + auto& queued = queued_it->second; auto* batch = queued.batch; auto& task = batch->task_list[queued.owner_task_id]; task.dispatch_time = std::chrono::steady_clock::now(); @@ -2629,13 +3189,43 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (staging_params.empty()) findStagingPolicy(task.request, staging_params); if (!staging_params.empty() && staging_proxy_) { + bool acquired = true; + auto credit_status = Status::OK(); + if (!queued.receiver_credit_exempt) { + acquired = false; + credit_status = acquireReceiverCredit(task.request, + queued.credit, acquired); + } + if (!credit_status.ok()) { + LOG(WARNING) << "Receiver-credit gate rejected staging " + "dispatch: " + << credit_status.ToString(); + return finishQueuedOwner(owner_id, FAILED); + } + if (!acquired) { + task.type = UNSPEC; + task.sub_task_id = -1; + CHECK_STATUS(runtime_queue_->defer(owner_id)); + return Status::OK(); + } + if (queued.credit.active) { + task.request.receiver_credit_session_high = + queued.credit.key.receiver_session.high; + task.request.receiver_credit_session_low = + queued.credit.key.receiver_session.low; + task.request.receiver_credit_epoch = queued.credit.epoch; + } task.staging = true; // Orchestration only; the real transport submissions issued by // ProxyManager are counted where they recurse through the // non-staging path below. auto status = staging_proxy_->submit(&task, (BatchID)batch, staging_params); - if (!status.ok()) return finishQueuedOwner(owner_id, FAILED); + if (!status.ok()) { + rollbackReceiverCredit(queued.credit); + return finishQueuedOwner(owner_id, FAILED); + } + if (queued.credit.active) queued.credit.committed = true; task.post_time = std::chrono::steady_clock::now(); return markQueuedOwnerSubmitted(owner_id); } @@ -2657,6 +3247,30 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { sub_batch->device_mask = task.device_mask; if (task.type == RDMA) sub_batch->qp_pool = task.qp_pool; } + + bool acquired = true; + auto credit_status = Status::OK(); + if (!queued.receiver_credit_exempt) { + acquired = false; + credit_status = + acquireReceiverCredit(task.request, queued.credit, acquired); + } + if (!credit_status.ok()) { + LOG(WARNING) << "Receiver-credit gate rejected dispatch: " + << credit_status.ToString(); + return finishQueuedOwner(owner_id, FAILED); + } + if (!acquired) { + task.type = UNSPEC; + task.sub_task_id = -1; + CHECK_STATUS(runtime_queue_->defer(owner_id)); + return Status::OK(); + } + if (queued.receiver_credit_exempt) { + bindReceiverCreditFence(sub_batch, task.request); + } else { + bindReceiverCreditFence(sub_batch, queued.credit); + } task.sub_task_id = sub_batch->size(); startTransportAttempt(task, task.type, std::chrono::steady_clock::now()); auto status = transport->submitTransferTasks(sub_batch, {task.request}); @@ -2669,8 +3283,12 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { if (attemptSubmitStageFailover(batch, queued.owner_task_id)) { return markQueuedOwnerSubmitted(owner_id); } + rollbackReceiverCredit(queued.credit); + task.type = UNSPEC; return finishQueuedOwner(owner_id, FAILED); } + if (queued.credit.active) queued.credit.committed = true; + task.post_time = std::chrono::steady_clock::now(); return markQueuedOwnerSubmitted(owner_id); } @@ -2712,6 +3330,7 @@ Status TransferEngineImpl::progressRuntimeQueue() { std::lock_guard lk(progress_mutex_); if (!runtime_queue_config_.enabled) return Status::OK(); + progressReceiverCreditReleases(); CHECK_STATUS(refillDispatchWindow()); std::vector owner_ids; @@ -2756,7 +3375,8 @@ Status TransferEngineImpl::progressRuntimeQueue() { bool TransferEngineImpl::hasActiveRuntimeQueue() { std::lock_guard lk(progress_mutex_); - return runtime_queue_config_.enabled && !queued_owners_.empty(); + return runtime_queue_config_.enabled && + (!queued_owners_.empty() || !pending_credit_releases_.empty()); } bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared, @@ -2764,6 +3384,12 @@ bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared, if (!runtime_queue_config_.enabled) return false; if (prepared.require_post_submit_cancellation) return false; if (owner_kind == QueueOwnerKind::StagingInternal) return true; + if (std::any_of(prepared.owners.begin(), prepared.owners.end(), + [](const PreparedSubmit::Owner& owner) { + return owner.receiver_credit_required; + })) { + return true; + } return std::none_of( prepared.owners.begin(), prepared.owners.end(), [](const PreparedSubmit::Owner& owner) { return owner.staging; }); @@ -2778,8 +3404,8 @@ Status TransferEngineImpl::submitTransfer( BatchRef batch_ref(*this, batch); const size_t start_task_id = batch_ref.get()->task_list.size(); PreparedSubmit prepared; - CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, prepared, - require_post_submit_cancellation)); + CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, owner_kind, + prepared, require_post_submit_cancellation)); if (shouldQueueSubmit(prepared, owner_kind)) { CHECK_STATUS( @@ -2967,6 +3593,18 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { auto& task = batch->task_list[task_id]; auto prev_type = task.type; + QueueOwnerId queue_owner_id = 0; + QueuedOwnerState* queued_owner = nullptr; + if (runtime_queue_config_.enabled && batch->queue_token != 0 && + runtime_queue_ + ->resolveOwner(batch->queue_token, task_id, queue_owner_id) + .ok()) { + auto queued_it = queued_owners_.find(queue_owner_id); + if (queued_it != queued_owners_.end()) { + queued_owner = &queued_it->second; + } + } + if (++task.failover_count > task.runtime_policy.max_failover_attempts) { LOG(WARNING) << "Task failover limit reached (" << task.runtime_policy.max_failover_attempts @@ -3027,13 +3665,72 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { } task.sub_task_id = sub_batch->size(); task.type = type; + ReceiverCreditTarget credit_target; + const bool credit_exempt = + queued_owner && queued_owner->receiver_credit_exempt; + if (!credit_exempt) { + CHECK_STATUS(resolveReceiverCreditTarget(task.request, credit_target)); + if (credit_target.required && !queued_owner) { + return Status::InvalidEntry( + "credit-required failover needs runtime queue " + "ownership" LOC_MARK); + } + } + if (queued_owner && queued_owner->credit.active && + (!credit_target.required || + queued_owner->credit.key.receiver_session != + credit_target.advert.receiver_session_id || + queued_owner->credit.epoch != credit_target.advert.epoch)) { + // A reservation from another receiver incarnation cannot authorize + // this attempt. Do not report it to a possibly replaced authority; + // abandon it and acquire against the current session below. + queued_owner->credit = ReceiverCreditReservation{}; + } + bool acquired_this_attempt = false; + if (!credit_exempt && queued_owner && !queued_owner->credit.active) { + bool acquired = false; + CHECK_STATUS(acquireReceiverCredit(task.request, queued_owner->credit, + acquired)); + if (!acquired) { + if (!queued_owner->in_dispatch_window || + dispatch_inflight_owners_ == 0 || + dispatch_inflight_bytes_ < queued_owner->byte_charge) { + return Status::InternalError( + "runtime dispatch window accounting underflow" LOC_MARK); + } + --dispatch_inflight_owners_; + dispatch_inflight_bytes_ -= queued_owner->byte_charge; + queued_owner->in_dispatch_window = false; + CHECK_STATUS(runtime_queue_->defer(queue_owner_id)); + task.type = UNSPEC; + task.sub_task_id = -1; + return Status::OK(); + } + acquired_this_attempt = queued_owner->credit.active; + } + if (queued_owner) { + if (queued_owner->receiver_credit_exempt) { + bindReceiverCreditFence(sub_batch, task.request); + } else { + bindReceiverCreditFence(sub_batch, queued_owner->credit); + } + } startTransportAttempt(task, type, std::chrono::steady_clock::now()); - auto status = transport->submitTransferTasks(sub_batch, {task.request}); - if (!status.ok()) { + auto submit_status = + transport->submitTransferTasks(sub_batch, {task.request}); + if (!submit_status.ok()) { if (task.failure_stage < 0) task.failure_stage = 0; finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + if (queued_owner && acquired_this_attempt) { + rollbackReceiverCredit(queued_owner->credit); + } + return submit_status; } - return status; + if (queued_owner && queued_owner->credit.active) { + queued_owner->credit.committed = true; + } + task.post_time = std::chrono::steady_clock::now(); + return Status::OK(); } bool TransferEngineImpl::attemptSubmitStageFailover(Batch* batch, @@ -3641,9 +4338,32 @@ Status TransferEngineImpl::waitTransferCompletion(BatchID batch_id) { Status TransferEngineImpl::transferSync( const std::vector& request_list) { + auto has_any_fence = [](const Request& request) { + return request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0 || + request.receiver_credit_epoch != 0; + }; + auto has_complete_fence = [](const Request& request) { + return (request.receiver_credit_session_high != 0 || + request.receiver_credit_session_low != 0) && + request.receiver_credit_epoch != 0; + }; + const bool delegated_credit_fence = + !request_list.empty() && + std::all_of(request_list.begin(), request_list.end(), + has_complete_fence); + if (!delegated_credit_fence && + std::any_of(request_list.begin(), request_list.end(), has_any_fence)) { + return Status::InvalidArgument( + "mixed or partial delegated receiver-credit fence" LOC_MARK); + } + auto batch_id = allocateBatch(request_list.size()); BatchGuard batch_guard(*this, batch_id); - CHECK_STATUS(submitTransfer(batch_id, request_list)); + CHECK_STATUS(submitTransfer(batch_id, request_list, nullptr, + delegated_credit_fence + ? QueueOwnerKind::StagingInternal + : QueueOwnerKind::User)); for (uint64_t poll_count = 0;; ++poll_count) { TransferStatus xfer_status; CHECK_STATUS(progressBatch(batch_id, xfer_status)); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index 4b64333b6c..0331a05cba 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -540,6 +540,11 @@ Status RdmaTransport::submitTransferTasks( task->request = request; task->device_mask = rdma_batch->device_mask; task->qp_pool = rdma_batch->qp_pool; // RFC #2568 step 3 + task->receiver_credit_session_high = + rdma_batch->receiver_credit_session_high; + task->receiver_credit_session_low = + rdma_batch->receiver_credit_session_low; + task->receiver_credit_epoch = rdma_batch->receiver_credit_epoch; task->num_slices = 0; task->status_word = PENDING; task->transferred_bytes = 0; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 468cddc2b2..0914c05dd8 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -1835,6 +1835,24 @@ Status Workers::generatePostPath(RdmaSlice* slice) { auto target_id = slice->task->request.target_id; CHECK_STATUS(getRouteHint(target, target_id, (uint64_t)slice->target_addr, slice->length)); + const bool has_credit_fence = + slice->task->receiver_credit_session_high != 0 || + slice->task->receiver_credit_session_low != 0 || + slice->task->receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = target.segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + slice->task->receiver_credit_session_high || + advert->receiver_session_id.low != + slice->task->receiver_credit_session_low || + advert->epoch != slice->task->receiver_credit_epoch) { + return Status::InvalidEntry( + "RDMA receiver-credit session changed before post" LOC_MARK); + } + } if (slice->retry_count == 0) CHECK_STATUS(selectOptimalDevice(source, target, slice)); diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp index ef6162dd07..d7288c5376 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp @@ -176,6 +176,9 @@ Status TcpTransport::submitTransferTasks( task.request = request; task.progress_batch_id = batch->progress_batch_id; task.notify_progress = batch->notify_progress; + task.receiver_credit_session_high = batch->receiver_credit_session_high; + task.receiver_credit_session_low = batch->receiver_credit_session_low; + task.receiver_credit_epoch = batch->receiver_credit_epoch; task.status_word.store(TransferStatusEnum::PENDING, std::memory_order_release); } @@ -262,9 +265,7 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { return Status::InternalError("Transport shutting down"); std::string rpc_server_addr; - auto status = - findRemoteSegment(task->request.target_offset, task->request.length, - task->request.target_id, rpc_server_addr); + auto status = findRemoteSegment(*task, rpc_server_addr); if (!status.ok()) return status; Status last_error; @@ -294,11 +295,15 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { if (task->request.opcode == Request::WRITE) { status = ControlClient::sendData( rpc_server_addr, task->request.target_offset, - task->request.source, task->request.length); + task->request.source, task->request.length, + task->receiver_credit_session_high, + task->receiver_credit_session_low, task->receiver_credit_epoch); } else { status = ControlClient::recvData( rpc_server_addr, task->request.target_offset, - task->request.source, task->request.length); + task->request.source, task->request.length, + task->receiver_credit_session_high, + task->receiver_credit_session_low, task->receiver_credit_epoch); } if (status.ok()) return Status::OK(); @@ -325,9 +330,7 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { CHECK_STATUS(metadata_->segmentManager().invalidateRemote( task->request.target_id)); if (attempt == params_.max_retry_count) break; - auto resolve = findRemoteSegment( - task->request.target_offset, task->request.length, - task->request.target_id, rpc_server_addr); + auto resolve = findRemoteSegment(*task, rpc_server_addr); if (!resolve.ok()) return resolve; } } @@ -335,12 +338,35 @@ Status TcpTransport::doTransferWithRetry(TcpTask *task) { return last_error; } -Status TcpTransport::findRemoteSegment(uint64_t dest_addr, uint64_t length, - uint64_t target_id, +Status TcpTransport::findRemoteSegment(const TcpTask &task, std::string &rpc_server_addr) { return metadata_->segmentManager().withCachedSegment( - target_id, [&](SegmentDesc *segment) { - auto buffer = segment->findBuffer(dest_addr, length); + task.request.target_id, [&](SegmentDesc *segment) { + const bool has_credit_fence = + task.receiver_credit_session_high != 0 || + task.receiver_credit_session_low != 0 || + task.receiver_credit_epoch != 0; + if (has_credit_fence) { + if (segment->type != SegmentType::Memory) { + return Status::InvalidMetadataType( + "TCP receiver-credit target is not memory" LOC_MARK); + } + const auto &advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + task.receiver_credit_session_high || + advert->receiver_session_id.low != + task.receiver_credit_session_low || + advert->epoch != task.receiver_credit_epoch) { + return Status::InvalidEntry( + "TCP receiver-credit session changed before " + "send" LOC_MARK); + } + } + auto buffer = segment->findBuffer(task.request.target_offset, + task.request.length); rpc_server_addr = segment->rpc_server_addr; if (!buffer) { return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp index 665e9204d1..e887bbf158 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/endpoint.cpp @@ -34,7 +34,11 @@ bool samePeer(const UbBootstrapDesc& lhs, const UbBootstrapDesc& rhs) { return lhs.protocol_version == rhs.protocol_version && lhs.local_eid == rhs.local_eid && lhs.jetty_ids == rhs.jetty_ids && lhs.jetty_uasids == rhs.jetty_uasids && - lhs.endpoint_generation == rhs.endpoint_generation; + lhs.endpoint_generation == rhs.endpoint_generation && + lhs.receiver_credit_session_high == + rhs.receiver_credit_session_high && + lhs.receiver_credit_session_low == rhs.receiver_credit_session_low && + lhs.receiver_credit_epoch == rhs.receiver_credit_epoch; } } // namespace @@ -340,10 +344,26 @@ Status UbEndpoint::makeBootstrapDesc(const std::string& segment_name, } output.endpoint_generation = generation_; output.segment_generation = segment_generation; - output.capabilities = {"read", "write", "endpoint_generation"}; + output.capabilities = {"read", "write", "endpoint_generation", + "receiver_credit_v1"}; return Status::OK(); } +bool UbEndpoint::peerSupportsCapability(std::string_view capability) const { + std::lock_guard lock(lifecycle_mutex_); + return std::find(peer_.capabilities.begin(), peer_.capabilities.end(), + capability) != peer_.capabilities.end(); +} + +bool UbEndpoint::peerMatchesReceiverCredit(uint64_t session_high, + uint64_t session_low, + uint64_t epoch) const { + std::lock_guard lock(lifecycle_mutex_); + return peer_.receiver_credit_session_high == session_high && + peer_.receiver_credit_session_low == session_low && + peer_.receiver_credit_epoch == epoch; +} + bool UbEndpoint::tryAcquireOutstanding(uint64_t bytes) noexcept { std::lock_guard lock(lifecycle_mutex_); if (state_.load(std::memory_order_relaxed) != State::kReady || !context_ || diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp index 652b18c6d3..59c9c14422 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -36,6 +37,8 @@ namespace mooncake::tent { namespace { +constexpr std::string_view kReceiverCreditCapability = "receiver_credit_v1"; + bool filterAllows(const std::vector& filter, const ub::DeviceInfo& device) { if (filter.empty()) return true; @@ -45,6 +48,43 @@ bool filterAllows(const std::vector& filter, filter.end(); } +bool receiverCreditRequired(const SegmentDesc* segment) { + if (!segment || segment->type != SegmentType::Memory) return false; + const auto& advert = segment->getMemory().receiver_credit; + return advert && advert->schema_version == kReceiverCreditProtocolVersion && + advert->flags == kReceiverCreditRequired && + !advert->receiver_session_id.empty() && advert->epoch != 0; +} + +const ReceiverCreditAdvertV1* receiverCreditAdvert(const SegmentDesc* segment) { + if (!receiverCreditRequired(segment)) return nullptr; + return &*segment->getMemory().receiver_credit; +} + +void bindReceiverCreditIdentity(const SegmentDesc* segment, + UbBootstrapDesc& bootstrap) { + const auto* advert = receiverCreditAdvert(segment); + if (!advert) return; + bootstrap.receiver_credit_session_high = advert->receiver_session_id.high; + bootstrap.receiver_credit_session_low = advert->receiver_session_id.low; + bootstrap.receiver_credit_epoch = advert->epoch; +} + +bool matchesReceiverCredit(const UbBootstrapDesc& bootstrap, + const ReceiverCreditAdvertV1& advert) { + return bootstrap.receiver_credit_session_high == + advert.receiver_session_id.high && + bootstrap.receiver_credit_session_low == + advert.receiver_session_id.low && + bootstrap.receiver_credit_epoch == advert.epoch; +} + +bool supportsReceiverCredit(const UbBootstrapDesc& bootstrap) { + return std::find(bootstrap.capabilities.begin(), + bootstrap.capabilities.end(), + kReceiverCreditCapability) != bootstrap.capabilities.end(); +} + uint64_t generationSeed() { const auto wall = static_cast( std::chrono::duration_cast( @@ -510,7 +550,22 @@ struct UbTransport::Impl { request.remote_topology_id, peer_path}; CHECK_STATUS( endpoints->getOrCreate(key, request.local_context, endpoint)); + const auto* credit_advert = + receiverCreditAdvert(request.remote_segment); + const bool credit_required = credit_advert != nullptr; if (endpoint->ready()) { + if (credit_required && + (!endpoint->peerSupportsCapability(kReceiverCreditCapability) || + !endpoint->peerMatchesReceiverCredit( + credit_advert->receiver_session_id.high, + credit_advert->receiver_session_id.low, + credit_advert->epoch))) { + (void)endpoints->retire(endpoint); + endpoint.reset(); + return Status::InvalidArgument( + "Remote UB endpoint lacks required receiver-credit " + "support" LOC_MARK); + } recordReadyEndpoint(endpoint); return Status::OK(); } @@ -527,9 +582,20 @@ struct UbTransport::Impl { endpoint.reset(); return status; } + auto local_segment = + metadata ? metadata->segmentManager().getLocal() : SegmentDescRef{}; + bindReceiverCreditIdentity(local_segment.get(), bootstrap); UbBootstrapDesc response; status = ControlClient::bootstrapUb( request.remote_segment->rpc_server_addr, bootstrap, response); + if (status.ok() && credit_required) { + if (!supportsReceiverCredit(response) || + !matchesReceiverCredit(response, *credit_advert)) { + status = Status::InvalidArgument( + "Remote UB bootstrap receiver-credit identity " + "mismatch" LOC_MARK); + } + } if (status.ok()) status = endpoint->bind(response); if (!status.ok()) { (void)endpoints->retire(endpoint); @@ -545,6 +611,14 @@ struct UbTransport::Impl { response.reply_msg = "UB transport is shutting down"; return -1; } + auto local_segment = + metadata ? metadata->segmentManager().getLocal() : SegmentDescRef{}; + if (receiverCreditRequired(local_segment.get()) && + !supportsReceiverCredit(request)) { + response.reply_msg = + "UB bootstrap peer lacks required receiver-credit support"; + return -1; + } const std::string local_name = getNicNameFromNicPath(request.peer_nic_path); auto context_it = context_by_topology_name.find(local_name); @@ -571,6 +645,7 @@ struct UbTransport::Impl { request.local_nic_path, currentSegmentGeneration(), response); } if (status.ok()) { + bindReceiverCreditIdentity(local_segment.get(), response); recordReadyEndpoint(endpoint); } if (!status.ok()) { @@ -705,6 +780,7 @@ Status UbTransport::submitTransferTasks( const auto notify_progress = ub_batch->notify_progress; const auto progress_batch_id = ub_batch->progress_batch_id; + auto progress_armed = std::make_shared>(false); std::vector new_tasks; new_tasks.reserve(request_list.size()); for (const auto& request : request_list) { @@ -717,9 +793,12 @@ Status UbTransport::submitTransferTasks( "transports/ub/max_slices_per_task" LOC_MARK); } auto task = ub::UbTask::create( - request, - [notify_progress, progress_batch_id](const TransferStatus&) { - if (notify_progress) notify_progress(progress_batch_id); + request, [notify_progress, progress_batch_id, + progress_armed](const TransferStatus&) { + if (progress_armed->load(std::memory_order_acquire) && + notify_progress) { + notify_progress(progress_batch_id); + } }); size_t offset = 0; while (offset < request.length) { @@ -731,6 +810,11 @@ Status UbTransport::submitTransferTasks( spec.length = length; spec.request_offset = offset; spec.max_retries = impl_->params.max_retries; + spec.receiver_credit_session_high = + ub_batch->receiver_credit_session_high; + spec.receiver_credit_session_low = + ub_batch->receiver_credit_session_low; + spec.receiver_credit_epoch = ub_batch->receiver_credit_epoch; if (!task->addSlice(spec)) { return Status::InternalError( "Unable to construct UB slice" LOC_MARK); @@ -741,15 +825,15 @@ Status UbTransport::submitTransferTasks( new_tasks.push_back(std::move(task)); } - for (auto& task : new_tasks) { - ub_batch->task_list.push_back(task); - auto status = impl_->workers->submit(task, ub_batch->device_mask); - if (!status.ok()) { - for (auto& queued : new_tasks) { - if (queued) queued->requestCancellation(); - } - return status; - } + CHECK_STATUS(impl_->workers->submitBatch(new_tasks, ub_batch->device_mask)); + ub_batch->task_list.insert(ub_batch->task_list.end(), new_tasks.begin(), + new_tasks.end()); + progress_armed->store(true, std::memory_order_release); + // A very small task can finish between worker publication and the batch + // append above. Its callback intentionally stays disarmed until the task + // is queryable; this notification closes that race. + if (!new_tasks.empty() && notify_progress) { + notify_progress(progress_batch_id); } return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp index effe77fff2..28e3dc0e79 100644 --- a/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ub/workers.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -275,44 +276,101 @@ Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) { if (!task) { return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK); } + return submitBatch({task}, device_mask); +} + +Status UbWorkers::submitBatch(const std::vector& tasks, + uint64_t device_mask) { if (!accepting_.load(std::memory_order_acquire)) { return Status::InvalidArgument( "UB workers are not accepting work" LOC_MARK); } - const auto snapshot = task->snapshot(); - if (!snapshot.sealed) { - return Status::InvalidArgument( - "UB task must be sealed before submission" LOC_MARK); - } - const int priority = - std::clamp(task->request().priority, static_cast(PRIO_HIGH), - static_cast(PRIO_LOW)); std::vector pending; - for (const auto& slice : task->slices()) { - if (slice && slice->markQueued()) { + std::unordered_set unique_tasks; + for (const auto& task : tasks) { + if (!task) { + return Status::InvalidArgument( + "Cannot submit a null UB task" LOC_MARK); + } + if (!unique_tasks.insert(task.get()).second) { + return Status::InvalidArgument( + "Cannot submit the same UB task twice in one batch" LOC_MARK); + } + const auto snapshot = task->snapshot(); + if (!snapshot.sealed) { + return Status::InvalidArgument( + "UB task must be sealed before submission" LOC_MARK); + } + const int priority = + std::clamp(task->request().priority, static_cast(PRIO_HIGH), + static_cast(PRIO_LOW)); + for (const auto& slice : task->slices()) { + if (!slice || slice->snapshot().state != UbSliceState::kInitial) { + return Status::InvalidArgument( + "UB task contains a slice that cannot be " + "submitted" LOC_MARK); + } pending.push_back(PendingSlice{task, slice, device_mask, priority, task->request().target_id, task->request().opcode}); } } - bool stopped_during_submit = false; + + Status failure = Status::OK(); + bool notify = false; + size_t marked_count = 0; { std::lock_guard lock(queue_mutex_); if (!accepting_.load(std::memory_order_relaxed)) { - stopped_during_submit = true; - } else { - for (auto& item : pending) { - queues_[item.priority].push_back(std::move(item)); + failure = Status::InvalidArgument( + "UB workers stopped during submission" LOC_MARK); + } + + if (failure.ok()) { + for (const auto& item : pending) { + if (!item.slice->markQueued()) { + failure = Status::InvalidArgument( + "UB task changed while it was being " + "submitted" LOC_MARK); + break; + } + ++marked_count; + } + } + + if (failure.ok()) { + std::array inserted_counts{}; + try { + for (const auto& item : pending) { + queues_[item.priority].push_back(item); + ++inserted_counts[item.priority]; + } + notify = !pending.empty(); + } catch (const std::bad_alloc&) { + for (size_t priority = 0; priority < queues_.size(); + ++priority) { + while (inserted_counts[priority] != 0) { + queues_[priority].pop_back(); + --inserted_counts[priority]; + } + } + failure = Status::InternalError( + "Unable to enqueue UB task batch" LOC_MARK); } } } - if (stopped_during_submit) { - for (auto& item : pending) item.slice->requestCancellation(); - return Status::InvalidArgument( - "UB workers stopped during submission" LOC_MARK); + + if (!failure.ok()) { + // A markQueued() race or allocation failure may occur after earlier + // slices changed state. None of them are visible to posting workers; + // cancel the private tasks so no slice remains orphaned in kQueued. + for (size_t index = 0; index < marked_count; ++index) { + (void)pending[index].task->requestCancellation(); + } + return failure; } - if (!pending.empty()) queue_cv_.notify_all(); + if (notify) queue_cv_.notify_all(); return Status::OK(); } @@ -472,9 +530,28 @@ Status UbWorkers::buildRoute(const PendingSlice& pending, Route& route) { return Status::InvalidMetadataType( "UB target is not a memory segment" LOC_MARK); } + const auto& spec = pending.slice->spec(); + const bool has_credit_fence = + spec.receiver_credit_session_high != 0 || + spec.receiver_credit_session_low != 0 || + spec.receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + spec.receiver_credit_session_high || + advert->receiver_session_id.low != + spec.receiver_credit_session_low || + advert->epoch != spec.receiver_credit_epoch) { + return Status::InvalidEntry( + "UB receiver-credit session changed before " + "post" LOC_MARK); + } + } auto* buffer = - segment->findBuffer(pending.slice->spec().remote_address, - pending.slice->spec().length); + segment->findBuffer(spec.remote_address, spec.length); if (!buffer) { return Status::NeedsRefreshCache( "UB target range is not registered" LOC_MARK); diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index ab404fff3c..438bb3057e 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -252,6 +252,193 @@ TEST(AdmissionQueueTest, KeepsAdmissionOrderForDispatch) { EXPECT_EQ(picked, expected_ids); } +TEST(AdmissionQueueTest, DefersDispatchingOwnerToFifoTailWithoutReaccounting) { + LocalTransferAdmissionQueue queue({4, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {makeOwner(0, 16), makeOwner(1, 24), + makeOwner(2, 32)}), + admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 3u); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + const size_t owners_before = queue.outstandingOwners(); + const size_t bytes_before = queue.outstandingBytes(); + + EXPECT_TRUE(queue.defer(picked[0]).ok()); + EXPECT_EQ(queue.outstandingOwners(), owners_before); + EXPECT_EQ(queue.outstandingBytes(), bytes_before); + + picked = queue.pickForDispatch(3, 128); + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfOwnerYieldsToCurrentWaiters) { + QueueLimits limits{4, 128, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 100; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 200; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + EXPECT_TRUE(queue.defer(picked[0]).ok()); + + picked = queue.pickForDispatch(3, 128); + // A temporarily blocked earliest-deadline owner yields one round so it + // cannot starve later-deadline work that may have available resources. + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfOwnerKeepsRoundAheadOfNewAdmission) { + QueueLimits limits{4, 128, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 200; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 300; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto newcomer = makeOwner(0, 16); + newcomer.request.deadline_ns = 150; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(2, 1, {std::move(newcomer)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids, std::vector({4})); + + picked = queue.pickForDispatch(4, 128); + // The deferred owner starts the next EDF round. A newly admitted owner is + // ordered within that round instead of jumping ahead of owners 2 and 3, + // which were already waiting when owner 1 yielded. + const std::vector expected{2, 3, 1, 4}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeferredEdfRoundsStayOrderedAcrossAdmissions) { + QueueLimits limits{6, 256, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + auto first = makeOwner(0, 16); + first.request.deadline_ns = 100; + auto second = makeOwner(1, 16); + second.request.deadline_ns = 200; + auto third = makeOwner(2, 16); + third.request.deadline_ns = 300; + ASSERT_TRUE(queue + .tryAdmit(makeSubmit(1, 3, + {std::move(first), std::move(second), + std::move(third)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 256); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto deadline_250 = makeOwner(0, 16); + deadline_250.request.deadline_ns = 250; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(2, 1, {std::move(deadline_250)}), admitted_ids) + .ok()); + + picked = queue.pickForDispatch(1, 256); + ASSERT_EQ(picked, std::vector({2})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + auto deadline_225 = makeOwner(0, 16); + deadline_225.request.deadline_ns = 225; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(3, 1, {std::move(deadline_225)}), admitted_ids) + .ok()); + + picked = queue.pickForDispatch(5, 256); + const std::vector expected{3, 1, 2, 5, 4}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionCannotUndoCreditDeferralRound) { + QueueLimits limits{2, 128, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = 10; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{100}; }); + + auto starved = makeOwner(0, 16); + starved.request.deadline_ns = 105; + auto waiter = makeOwner(1, 16); + waiter.request.deadline_ns = 1000; + std::vector admitted_ids; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(1, 2, {std::move(starved), std::move(waiter)}), + admitted_ids) + .ok()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked, std::vector({1})); + ASSERT_TRUE(queue.defer(picked[0]).ok()); + + // Owner 1 remains urgent, but it yielded into round 1. Promotion cannot + // move it ahead of owner 2, which is still waiting in round 0. + picked = queue.pickForDispatch(1, 128); + EXPECT_EQ(picked, std::vector({2})); +} + +TEST(AdmissionQueueTest, RejectsDeferForNonDispatchingOwner) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 1u); + + EXPECT_TRUE(queue.defer(0).IsInvalidArgument()); + EXPECT_TRUE(queue.defer(99).IsInvalidEntry()); + EXPECT_TRUE(queue.defer(admitted_ids[0]).IsInvalidEntry()); + + auto picked = queue.pickForDispatch(1, 128); + ASSERT_EQ(picked.size(), 1u); + ASSERT_TRUE(queue.complete(picked[0], COMPLETED).ok()); + EXPECT_TRUE(queue.defer(picked[0]).IsInvalidEntry()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); +} + TEST(AdmissionQueueTest, RequiresDispatchBeforeTerminalCompletion) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; @@ -828,6 +1015,29 @@ TEST(AdmissionQueueTest, EXPECT_EQ(queue.dispatchingBytes(), 0u); } +TEST(AdmissionQueueTest, RejectsDeferForOwnerDroppedDuringDispatchPick) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + ASSERT_TRUE( + queue + .tryAdmit(makeSubmit(1, 1, + {makeDegradationEligibleOwnerWithDeadline( + 0, 16, 1'000'000'010)}), + admitted_ids) + .ok()); + + std::vector dropped; + EXPECT_TRUE(queue.pickForDispatch(1, 128, &dropped).empty()); + ASSERT_EQ(dropped, std::vector({1})); + EXPECT_TRUE(queue.defer(dropped[0]).IsInvalidEntry()); + EXPECT_TRUE(queue.pickForDispatch(1, 128).empty()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); +} + TEST(AdmissionQueueTest, Step3DropsAlreadyExpiredDeadline) { LocalTransferAdmissionQueue queue(step3Limits(1.5)); queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp index 33faa04a2c..86b74670a1 100644 --- a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp @@ -10,7 +10,7 @@ namespace mooncake::tent { namespace { -CreditKey key() { return {{1, 2}, 3, 4}; } +CreditKey key() { return {{1, 2}, {3, 4}, 1}; } ReceiverCreditUpdateV1 update(uint64_t epoch, uint64_t seq, std::vector grants) { ReceiverCreditUpdateV1 u; @@ -18,6 +18,7 @@ ReceiverCreditUpdateV1 update(uint64_t epoch, uint64_t seq, u.qos_class = key().qos_class; u.epoch = epoch; u.sequence = seq; + u.freshness_ttl_ms = 1000; u.grants = std::move(grants); return u; } @@ -126,7 +127,7 @@ TEST(ReceiverCredit, LedgerEntryCountIsBounded) { SenderCreditLedger l(1); ASSERT_TRUE(l.activate(key(), 7).ok()); auto other = key(); - ++other.sender_peer; + ++other.sender_instance.low; EXPECT_TRUE(l.activate(other, 7).IsTooManyRequests()); // A new epoch for an existing key does not consume another entry. EXPECT_TRUE(l.activate(key(), 8).ok()); @@ -138,7 +139,7 @@ TEST(ReceiverCredit, DeactivationReleasesCapacityAfterExactEpochFence) { grant(l, 1); ASSERT_TRUE(l.tryReserve(key(), charge(80, 1)).ok()); auto other = key(); - ++other.sender_peer; + ++other.sender_instance.low; EXPECT_TRUE(l.activate(other, 1).IsTooManyRequests()); EXPECT_TRUE(l.deactivate(key(), 6).IsInvalidEntry()); @@ -201,6 +202,27 @@ TEST(ReceiverCredit, RollbackChecksUnderflowAtomically) { ASSERT_TRUE(l.consumed(key(), CreditResource::RequestSlots, v).ok()); EXPECT_EQ(v, 1); ASSERT_TRUE(l.rollbackReservation(key(), charge(60, 1)).ok()); + ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 60); // cumulative consumption never moves backwards + ASSERT_TRUE(l.released(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 60); + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 40); // returned authorization is not locally re-minted +} + +TEST(ReceiverCredit, DeactivateSessionRemovesAllQosEntries) { + SenderCreditLedger ledger(2); + auto high = key(); + auto low = key(); + low.qos_class = 2; + ASSERT_TRUE(ledger.activate(high, 7).ok()); + ASSERT_TRUE(ledger.activate(low, 7).ok()); + ASSERT_TRUE(ledger.deactivateSession(key().receiver_session).ok()); + + auto replacement = key(); + replacement.receiver_session = {9, 10}; + replacement.sender_instance = {11, 12}; + ASSERT_TRUE(ledger.activate(replacement, 1).ok()); } TEST(ReceiverCredit, GrantCannotDecreaseOrFallBelowConsumption) { @@ -232,5 +254,335 @@ TEST(ReceiverCredit, ConcurrentReservationsNeverExceedGrant) { ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, consumed).ok()); EXPECT_EQ(consumed, 100); } + +TEST(ReceiverCredit, FreshnessExpiryFailsClosedAndReplayDoesNotRenew) { + uint64_t now_ns = 100; + SenderCreditLedger ledger(8, 1000, [&] { return now_ns; }); + ASSERT_TRUE(ledger.activate(key(), 7).ok()); + auto first = update( + 7, 1, + {{CreditResource::DataBytes, 100}, {CreditResource::RequestSlots, 2}}); + first.freshness_ttl_ms = 1; + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger.applyUpdate(key(), first, disposition).ok()); + now_ns += 999999; + EXPECT_TRUE(ledger.tryReserve(key(), charge(1, 1)).ok()); + + // A replay is accepted idempotently, but must not extend the lease. + ASSERT_TRUE(ledger.applyUpdate(key(), first, disposition).ok()); + EXPECT_EQ(disposition, CreditUpdateDisposition::DuplicateOrOld); + ++now_ns; + EXPECT_TRUE(ledger.tryReserve(key(), charge(1, 1)).IsInvalidEntry()); +} + +TEST(ReceiverCredit, TerminalReleaseRequiresReceiverReplenishment) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 2; + limits.freshness_ttl_ms = 1000; + ReceiverCreditAuthority authority({9, 10}, 3, limits); + ASSERT_TRUE(authority.status().ok()); + + CreditKey sender{{9, 10}, {11, 12}, 1}; + SenderCreditLedger ledger; + ASSERT_TRUE(ledger.activate(sender, 3).ok()); + + ReceiverCreditExchangeRequestV1 request; + ASSERT_TRUE(ledger.prepareExchange(sender, charge(60, 1), request).ok()); + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE(ledger.applyUpdate(sender, reply.update, disposition).ok()); + ASSERT_TRUE(ledger.tryReserve(sender, charge(60, 1)).ok()); + ASSERT_TRUE(ledger.releaseCommitted(sender, charge(60, 1)).ok()); + + uint64_t available = 0; + ASSERT_TRUE( + ledger.available(sender, CreditResource::DataBytes, available).ok()); + EXPECT_EQ(available, 0); // local release never mints credit + + ASSERT_TRUE(ledger.prepareExchange(sender, charge(60, 1), request).ok()); + ASSERT_TRUE(authority.exchange(request, reply).ok()); + ASSERT_TRUE(ledger.applyUpdate(sender, reply.update, disposition).ok()); + EXPECT_TRUE(ledger.tryReserve(sender, charge(60, 1)).ok()); + + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 60); +} + +TEST(ReceiverCredit, AuthorityGrantsMultiResourceDemandAtomically) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender, uint64_t bytes) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = { + {CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}, + }; + request.desired_grant_totals = { + {CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, 1}, + }; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 60), reply).ok()); + ASSERT_EQ(reply.update.grants.size(), kCreditResourceCount); + EXPECT_EQ(reply.update.grants[0].grant_total, 60); + EXPECT_EQ(reply.update.grants[1].grant_total, 1); + + ASSERT_TRUE(authority.exchange(make_request({5, 6}, 40), reply).ok()); + // The second sender cannot get a request slot, so bytes are not stranded + // in a partial grant either. + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + EXPECT_EQ(reply.update.grants[1].grant_total, 0); + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 60); +} + +TEST(ReceiverCredit, RollbackReleaseLetsAnotherSenderAcquireCapacity) { + ReceiverCreditLimits limits; + limits.capacity[0] = 60; + limits.capacity[1] = 1; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + SenderCreditLedger first; + CreditKey first_key{{1, 2}, {3, 4}, 0}; + ASSERT_TRUE(first.activate(first_key, 1).ok()); + + ReceiverCreditExchangeRequestV1 request; + ASSERT_TRUE(first.prepareExchange(first_key, charge(60, 1), request).ok()); + const auto initial_demand = request; + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + CreditUpdateDisposition disposition; + ASSERT_TRUE(first.applyUpdate(first_key, reply.update, disposition).ok()); + ASSERT_TRUE(first.tryReserve(first_key, charge(60, 1)).ok()); + ASSERT_TRUE(first.rollbackReservation(first_key, charge(60, 1)).ok()); + ASSERT_TRUE(first.prepareReleaseExchange(first_key, request).ok()); + const auto release_report = request; + ASSERT_TRUE(authority.exchange(request, reply).ok()); + EXPECT_EQ(reply.flags, 0); + + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 0); + + // Both a delayed initial demand and a replayed release receive retained + // cumulative state; neither can re-mint released capacity. + ASSERT_TRUE(authority.exchange(initial_demand, reply).ok()); + EXPECT_EQ(reply.flags, 0); + EXPECT_EQ(reply.update.grants[0].grant_total, 60); + ASSERT_TRUE(authority.exchange(release_report, reply).ok()); + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 0); + + SenderCreditLedger second; + CreditKey second_key{{1, 2}, {5, 6}, 0}; + ASSERT_TRUE(second.activate(second_key, 1).ok()); + ASSERT_TRUE( + second.prepareExchange(second_key, charge(60, 1), request).ok()); + ASSERT_TRUE(authority.exchange(request, reply).ok()); + ASSERT_TRUE(second.applyUpdate(second_key, reply.update, disposition).ok()); + EXPECT_TRUE(second.tryReserve(second_key, charge(60, 1)).ok()); +} + +TEST(ReceiverCredit, SenderLimitCountsSenderNotQosEntries) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 2; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender, uint32_t qos) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = qos; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 0), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({3, 4}, 2), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({5, 6}, 0), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + EXPECT_EQ(reply.update.grants[1].grant_total, 0); +} + +TEST(ReceiverCredit, ZeroGrantHistoryDoesNotConsumeActiveSenderLimit) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 1; + limits.max_entries = 4; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + auto first = make_request({3, 4}); + ReceiverCreditExchangeReplyV1 first_reply; + ASSERT_TRUE(authority.exchange(first, first_reply).ok()); + EXPECT_EQ(first_reply.update.grants[0].grant_total, 10); + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({5, 6}), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 0); + + first.known_update_sequence = first_reply.update.sequence; + first.report_sequence = 2; + first.released_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + ASSERT_TRUE(authority.exchange(first, reply).ok()); + + ASSERT_TRUE(authority.exchange(make_request({7, 8}), reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 10); + EXPECT_EQ(reply.update.grants[1].grant_total, 1); +} + +TEST(ReceiverCredit, ReplayHistoryLimitFailsClosed) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 1; + limits.max_entries = 2; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + auto make_request = [](SenderInstanceId sender) { + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = sender; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 10}, + {CreditResource::RequestSlots, 1}}; + return request; + }; + + ReceiverCreditExchangeReplyV1 reply; + ASSERT_TRUE(authority.exchange(make_request({3, 4}), reply).ok()); + ASSERT_TRUE(authority.exchange(make_request({5, 6}), reply).ok()); + EXPECT_TRUE( + authority.exchange(make_request({7, 8}), reply).IsTooManyRequests()); + + // Existing identities remain replayable after the cap is reached. + EXPECT_TRUE(authority.exchange(make_request({3, 4}), reply).ok()); +} + +TEST(ReceiverCredit, ReplayHistoryLimitMustCoverActiveSenders) { + ReceiverCreditLimits limits; + limits.capacity[0] = 10; + limits.capacity[1] = 1; + limits.max_senders = 2; + limits.max_entries = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + EXPECT_TRUE(authority.status().IsInvalidArgument()); +} + +TEST(ReceiverCredit, AdvertPublishesHardCapacityBounds) { + ReceiverCreditLimits limits; + limits.capacity[0] = 4096; + limits.capacity[1] = 8; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + const auto advert = authority.advert(); + ASSERT_EQ(advert.capacities.size(), 2u); + EXPECT_EQ(advert.capacities[0].resource, CreditResource::DataBytes); + EXPECT_EQ(advert.capacities[0].total, 4096u); + EXPECT_EQ(advert.capacities[1].resource, CreditResource::RequestSlots); + EXPECT_EQ(advert.capacities[1].total, 8u); + + const auto round_trip = + nlohmann::json(advert).get(); + EXPECT_EQ(round_trip.capacities.size(), advert.capacities.size()); + EXPECT_EQ(round_trip.capacities[0].total, 4096u); +} + +TEST(ReceiverCredit, ExchangeReplayIsIdempotent) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = {3, 4}; + request.qos_class = 0; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}, + {CreditResource::RequestSlots, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 50}, + {CreditResource::RequestSlots, 1}}; + + ReceiverCreditExchangeReplyV1 first, replay; + ASSERT_TRUE(authority.exchange(request, first).ok()); + ASSERT_TRUE(authority.exchange(request, replay).ok()); + EXPECT_EQ(replay.update.sequence, first.update.sequence); + EXPECT_EQ(replay.update.grants[0].grant_total, 50); + uint64_t outstanding = 0; + ASSERT_TRUE( + authority.outstanding(CreditResource::DataBytes, outstanding).ok()); + EXPECT_EQ(outstanding, 50); +} + +TEST(ReceiverCredit, InvalidNewSenderCannotConsumeAuthorityEntry) { + ReceiverCreditLimits limits; + limits.capacity[0] = 100; + limits.capacity[1] = 1; + limits.max_senders = 1; + ReceiverCreditAuthority authority({1, 2}, 1, limits); + + ReceiverCreditExchangeRequestV1 request; + request.expected_receiver_session_id = {1, 2}; + request.epoch = 1; + request.sender_instance_id = {3, 4}; + request.qos_class = 0; + request.known_update_sequence = 99; + request.report_sequence = 1; + request.released_totals = {{CreditResource::DataBytes, 0}}; + request.desired_grant_totals = {{CreditResource::DataBytes, 50}, + {CreditResource::RequestSlots, 1}}; + ReceiverCreditExchangeReplyV1 reply; + EXPECT_TRUE(authority.exchange(request, reply).IsInvalidArgument()); + + request.sender_instance_id = {5, 6}; + request.known_update_sequence = 0; + EXPECT_TRUE(authority.exchange(request, reply).ok()); + EXPECT_EQ(reply.update.grants[0].grant_total, 50); +} } // namespace } // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp index 116616e30e..638dd2ff7d 100644 --- a/mooncake-transfer-engine/tent/tests/request_merge_test.cpp +++ b/mooncake-transfer-engine/tent/tests/request_merge_test.cpp @@ -24,6 +24,7 @@ struct RequestBoundaryInfo { std::optional source_key; std::optional target_key; uint64_t max_merge_bytes{std::numeric_limits::max()}; + std::optional receiver_credit_data_capacity; }; struct MergeResult { diff --git a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp index 41149ed4e1..dbaddebfc5 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -37,6 +37,12 @@ namespace mooncake { // left the normal lifecycle still holds. class TransferEngineImplTestPeer { public: + static tent::Status submitRequiringPostSubmitCancellation( + tent::TransferEngineImpl& engine, tent::BatchID batch_id, + const std::vector& requests) { + return engine.submitTransferRequiringPostSubmitCancellation(batch_id, + requests); + } static size_t queuedOwners(tent::TransferEngineImpl& engine) { return engine.queued_owners_.size(); } @@ -101,6 +107,10 @@ class FakeTransport : public Transport { std::atomic submit_calls{0}; std::atomic status_calls{0}; std::atomic cancel_calls{0}; + std::atomic fail_next_submits{0}; + std::atomic last_credit_session_high{0}; + std::atomic last_credit_session_low{0}; + std::atomic last_credit_epoch{0}; bool cancellation_supported{true}; Status install(std::string&, std::shared_ptr, @@ -123,6 +133,21 @@ class FakeTransport : public Transport { Status submitTransferTasks(SubBatchRef batch, const std::vector& requests) override { ++submit_calls; + int failures = fail_next_submits.load(std::memory_order_relaxed); + while (failures > 0 && + !fail_next_submits.compare_exchange_weak( + failures, failures - 1, std::memory_order_relaxed)) { + } + if (failures > 0) { + return Status::InternalError( + "injected zero-handoff submit failure" LOC_MARK); + } + last_credit_session_high.store(batch->receiver_credit_session_high, + std::memory_order_relaxed); + last_credit_session_low.store(batch->receiver_credit_session_low, + std::memory_order_relaxed); + last_credit_epoch.store(batch->receiver_credit_epoch, + std::memory_order_relaxed); auto* fake = static_cast(batch); for (const auto& request : requests) { fake->requests.push_back(request); @@ -252,6 +277,13 @@ void installFakeRdma(TransferEngineImpl& engine, engine.swapTransportForTest(RDMA, fake_rdma); } +void installFakeUb(TransferEngineImpl& engine, + const std::shared_ptr& fake_ub) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake_ub->install(seg_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(UB, fake_ub); +} + Request makeLocalWrite(uint8_t* ptr, size_t length) { Request request; request.opcode = Request::WRITE; @@ -263,6 +295,26 @@ Request makeLocalWrite(uint8_t* ptr, size_t length) { return request; } +Request makeLocalUbWrite(uint8_t* ptr, size_t length) { + auto request = makeLocalWrite(ptr, length); + request.transport_hint = UB; + return request; +} + +Request makeRemoteUbWrite(uint8_t* source, SegmentID target_id, uint8_t* target, + size_t length) { + auto request = makeLocalUbWrite(source, length); + request.target_id = target_id; + request.target_offset = reinterpret_cast(target); + return request; +} + +Request makeLocalNetworkWrite(uint8_t* ptr, size_t length) { + auto request = makeLocalWrite(ptr, length); + request.transport_hint = UNSPEC; + return request; +} + TEST(RuntimeQueueDispatch, RejectsOverfullBatchBeforePublishingTasks) { auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); TransferEngineImpl engine(cfg); @@ -339,14 +391,13 @@ TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) { TransferEngineImpl engine(cfg); ASSERT_TRUE(engine.available()); - std::atomic complete_first{false}; + std::atomic complete{false}; auto fake_rdma = std::make_shared( - RDMA, [&complete_first](const Request& request, int) { - if (!complete_first.load()) { - return TransferStatus{TransferStatusEnum::PENDING, 0}; - } - return TransferStatus{TransferStatusEnum::COMPLETED, - request.length}; + RDMA, [&complete](const Request& request, int) { + return complete.load(std::memory_order_acquire) + ? TransferStatus{TransferStatusEnum::COMPLETED, + request.length} + : TransferStatus{TransferStatusEnum::PENDING, 0}; }); installFakeRdma(engine, fake_rdma); @@ -365,8 +416,7 @@ TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) { .ok()); EXPECT_EQ(fake_rdma->submit_calls.load(), 1); - complete_first.store(true); - + complete.store(true, std::memory_order_release); TransferStatus first{}; ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); @@ -1104,6 +1154,390 @@ TEST(RuntimeQueueDispatch, TcpCapsPreservePublicMappingAndOwnerBytes) { } } +TEST(RuntimeQueueDispatch, ScatterCannotBypassRequiredReceiverCredit) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_ub = std::make_shared(UB); + installFakeUb(engine, fake_ub); + std::vector buffer(kReqLen, 0x71); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + auto status = + TransferEngineImplTestPeer::submitRequiringPostSubmitCancellation( + engine, batch, {makeLocalUbWrite(buffer.data(), kReqLen)}); + EXPECT_TRUE(status.IsNotImplemented()); + EXPECT_EQ(fake_ub->submit_calls.load(), 0); + EXPECT_EQ(TransferEngineImplTestPeer::queuedOwners(engine), 0u); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ReceiverCreditDefersUbUntilTerminalRelease) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20, true); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + cfg->set("receiver_credit/freshness_ttl_ms", uint32_t{1000}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_ub = std::make_shared( + UB, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeUb(engine, fake_ub); + + std::vector buffer(kReqLen * 2, 0x71); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), kReqLen), + makeLocalUbWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + // Both owners fit the local dispatch window, but receiver credit permits + // only one physical UB request. Merging must stop at the advertised hard + // byte capacity, so the second owner remains queued instead of making an + // inadmissible 8 KiB owner. + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_ub->submit_calls.load(), 2); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ZeroHandoffUbFailureRollsCreditBack) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_ub = std::make_shared(UB); + fake_ub->fail_next_submits.store(1); + installFakeUb(engine, fake_ub); + std::vector buffer(kReqLen * 2, 0x72); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), kReqLen), + makeLocalUbWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + // The first submit failed before handoff. Its cumulative rollback is + // reported to the receiver before the next owner gets a fresh grant. + EXPECT_EQ(fake_ub->submit_calls.load(), 2); + TransferStatus first{}, second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::FAILED); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RemoteReceiverCreditExchangeGatesUbDispatch) { + constexpr size_t kReqLen = 4096; + auto receiver_cfg = makeRuntimeQueueConfig(2, 1UL << 20); + receiver_cfg->set("receiver_credit/receiver_enable", true); + receiver_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + receiver_cfg->set("receiver_credit/request_slots", uint64_t{1}); + auto sender_cfg = makeRuntimeQueueConfig(2, 1UL << 20); + sender_cfg->set("receiver_credit/sender_enable", true); + sender_cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + sender_cfg->set("receiver_credit/request_slots", uint64_t{1}); + + TransferEngineImpl receiver(receiver_cfg); + TransferEngineImpl sender(sender_cfg); + ASSERT_TRUE(receiver.available()); + ASSERT_TRUE(sender.available()); + + auto receiver_ub = std::make_shared(UB); + installFakeUb(receiver, receiver_ub); + std::atomic complete_first{false}; + auto sender_ub = std::make_shared( + UB, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeUb(sender, sender_ub); + + std::vector source(kReqLen * 2, 0x81); + std::vector target(kReqLen * 2, 0); + ASSERT_TRUE(sender.registerLocalMemory(source.data(), source.size()).ok()); + ASSERT_TRUE( + receiver.registerLocalMemory(target.data(), target.size()).ok()); + SegmentID remote = LOCAL_SEGMENT_ID; + ASSERT_TRUE(sender.openSegment(remote, receiver.getSegmentName()).ok()); + ASSERT_NE(remote, LOCAL_SEGMENT_ID); + + BatchID batch = sender.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + sender + .submitTransfer( + batch, {makeRemoteUbWrite(source.data(), remote, target.data(), + kReqLen), + makeRemoteUbWrite(source.data() + kReqLen, remote, + target.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(sender_ub->submit_calls.load(), 1); + + // TCP fallback carries the same receiver-session fence on the wire. The + // receiver rejects an unfenced or stale request even if the sender's + // metadata cache still describes the old process incarnation. + const auto receiver_addr = receiver.getSegmentName(); + const uint64_t session_high = sender_ub->last_credit_session_high.load(); + const uint64_t session_low = sender_ub->last_credit_session_low.load(); + const uint64_t epoch = sender_ub->last_credit_epoch.load(); + ASSERT_NE(session_high | session_low, 0u); + ASSERT_NE(epoch, 0u); + + auto delegated = makeLocalUbWrite(target.data(), kReqLen); + delegated.receiver_credit_session_high = session_high; + delegated.receiver_credit_session_low = session_low; + delegated.receiver_credit_epoch = epoch + 1; + EXPECT_TRUE( + ControlClient::delegate(receiver_addr, delegated).IsRpcServiceError()); + delegated.receiver_credit_epoch = epoch; + ASSERT_TRUE(ControlClient::delegate(receiver_addr, delegated).ok()); + EXPECT_EQ(receiver_ub->submit_calls.load(), 1); + + delegated.receiver_credit_epoch = 0; + EXPECT_TRUE( + ControlClient::delegate(receiver_addr, delegated).IsInvalidArgument()); + + EXPECT_TRUE(ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), + source.data(), kReqLen) + .IsRpcServiceError()); + EXPECT_EQ(target[0], 0); + EXPECT_TRUE(ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), + source.data(), kReqLen, session_high, session_low, + epoch + 1) + .IsRpcServiceError()); + EXPECT_EQ(target[0], 0); + const auto fenced_send = ControlClient::sendData( + receiver_addr, reinterpret_cast(target.data()), source.data(), + kReqLen, session_high, session_low, epoch); + ASSERT_TRUE(fenced_send.ok()) << fenced_send.ToString(); + EXPECT_EQ(target[0], source[0]); + + std::vector readback(kReqLen, 0); + EXPECT_TRUE(ControlClient::recvData( + receiver_addr, reinterpret_cast(target.data()), + readback.data(), kReqLen, session_high, session_low, + epoch + 1) + .IsRpcServiceError()); + const auto fenced_recv = ControlClient::recvData( + receiver_addr, reinterpret_cast(target.data()), + readback.data(), kReqLen, session_high, session_low, epoch); + ASSERT_TRUE(fenced_recv.ok()) << fenced_recv.ToString(); + EXPECT_EQ(readback, std::vector(kReqLen, 0x81)); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}, second{}; + ASSERT_TRUE(sender.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(sender_ub->submit_calls.load(), 2); + ASSERT_TRUE(sender.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(sender.freeBatch(batch).ok()); + EXPECT_TRUE(sender.closeSegment(remote).ok()); + EXPECT_TRUE( + sender.unregisterLocalMemory(source.data(), source.size()).ok()); + EXPECT_TRUE( + receiver.unregisterLocalMemory(target.data(), target.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RdmaFailureGatesUbFailoverWithReceiverCredit) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = + std::make_shared(RDMA, [](const Request&, int) { + return TransferStatus{TransferStatusEnum::FAILED, 0}; + }); + auto fake_ub = std::make_shared(UB); + installFakeRdma(engine, fake_rdma); + installFakeUb(engine, fake_ub); + + std::vector buffer(kReqLen, 0x91); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE(engine + .submitTransfer( + batch, {makeLocalNetworkWrite(buffer.data(), kReqLen)}) + .ok()); + ASSERT_EQ(fake_rdma->submit_calls.load(), 1); + ASSERT_EQ(fake_ub->submit_calls.load(), 0); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_ub->submit_calls.load(), 1); + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ReceiverCreditAlsoGatesInitialRdmaPath) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(2, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kReqLen}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_rdma = std::make_shared( + RDMA, [&complete_first](const Request& request, int) { + if (!complete_first.load(std::memory_order_acquire)) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + std::vector buffer(kReqLen * 2, 0x92); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_NE(fake_rdma->last_credit_session_high.load() | + fake_rdma->last_credit_session_low.load(), + 0u); + EXPECT_EQ(fake_rdma->last_credit_epoch.load(), 1u); + + complete_first.store(true, std::memory_order_release); + TransferStatus first{}, second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsRequestLargerThanReceiverCreditCapacity) { + constexpr size_t kCapacity = 4096; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("receiver_credit/enable", true); + cfg->set("receiver_credit/data_bytes", uint64_t{kCapacity}); + cfg->set("receiver_credit/request_slots", uint64_t{1}); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_ub = std::make_shared(UB); + installFakeUb(engine, fake_ub); + std::vector buffer(kCapacity * 2, 0x93); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + + const auto status = engine.submitTransfer( + batch, {makeLocalUbWrite(buffer.data(), buffer.size())}); + EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString(); + EXPECT_EQ(fake_ub->submit_calls.load(), 0); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, TcpControlDataUsesLegacyFramingWithoutCredit) { + constexpr size_t kLength = 256; + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + std::vector source(kLength, 0xa5); + std::vector target(kLength, 0); + std::vector readback(kLength, 0); + ASSERT_TRUE(engine.registerLocalMemory(target.data(), target.size()).ok()); + const auto server_addr = engine.getSegmentName(); + + ASSERT_TRUE(ControlClient::sendData( + server_addr, reinterpret_cast(target.data()), + source.data(), source.size()) + .ok()); + EXPECT_EQ(target, source); + ASSERT_TRUE(ControlClient::recvData( + server_addr, reinterpret_cast(target.data()), + readback.data(), readback.size()) + .ok()); + EXPECT_EQ(readback, source); + + EXPECT_TRUE(ControlClient::sendData( + server_addr, reinterpret_cast(target.data()), + source.data(), source.size(), 1, 0, 0) + .IsInvalidArgument()); + EXPECT_TRUE( + engine.unregisterLocalMemory(target.data(), target.size()).ok()); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index dffe707bc0..625133e097 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -411,6 +411,9 @@ TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { source.jetty_uasids = {21, 22}; source.endpoint_generation = 41; source.segment_generation = 42; + source.receiver_credit_session_high = 43; + source.receiver_credit_session_low = 44; + source.receiver_credit_epoch = 45; source.capabilities = {"read", "write"}; source.reply_msg = "ok"; @@ -427,6 +430,11 @@ TEST(ControlPlaneTest, UbBootstrapJsonRoundTripsNativeIdentity) { EXPECT_EQ(parsed.jetty_uasids, source.jetty_uasids); EXPECT_EQ(parsed.endpoint_generation, source.endpoint_generation); EXPECT_EQ(parsed.segment_generation, source.segment_generation); + EXPECT_EQ(parsed.receiver_credit_session_high, + source.receiver_credit_session_high); + EXPECT_EQ(parsed.receiver_credit_session_low, + source.receiver_credit_session_low); + EXPECT_EQ(parsed.receiver_credit_epoch, source.receiver_credit_epoch); EXPECT_EQ(parsed.capabilities, source.capabilities); EXPECT_EQ(parsed.reply_msg, source.reply_msg); } @@ -445,6 +453,9 @@ TEST(ControlPlaneTest, UbBootstrapJsonDefaultsOptionalFields) { EXPECT_TRUE(parsed.jetty_uasids.empty()); EXPECT_EQ(parsed.endpoint_generation, 0u); EXPECT_EQ(parsed.segment_generation, 0u); + EXPECT_EQ(parsed.receiver_credit_session_high, 0u); + EXPECT_EQ(parsed.receiver_credit_session_low, 0u); + EXPECT_EQ(parsed.receiver_credit_epoch, 0u); EXPECT_TRUE(parsed.capabilities.empty()); EXPECT_TRUE(parsed.reply_msg.empty()); } @@ -470,6 +481,7 @@ TEST(ControlPlaneTest, UbBootstrapRpcIdIsAppendedWithoutRenumbering) { EXPECT_EQ(static_cast(SubscribeSegmentUpdate), 10); EXPECT_EQ(static_cast(NotifySegmentUpdated), 11); EXPECT_EQ(static_cast(BootstrapUb), 12); + EXPECT_EQ(static_cast(ExchangeReceiverCredit), 13); } // --------------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp index 4dc2694de6..2ad2b862fe 100644 --- a/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp +++ b/mooncake-transfer-engine/tent/tests/ub_native_data_path_test.cpp @@ -3,6 +3,7 @@ #include +#include #include #include #include @@ -547,6 +548,30 @@ TEST(UbNativeDataPathTest, (void)endpoints.retire(endpoint); }); ASSERT_TRUE(workers.start().ok()); + + Request invalid_batch_request{}; + invalid_batch_request.opcode = Request::WRITE; + invalid_batch_request.source = source.data(); + invalid_batch_request.target_id = LOCAL_SEGMENT_ID; + invalid_batch_request.target_offset = + reinterpret_cast(target.data()); + invalid_batch_request.length = 16; + auto valid_task = UbTask::create(invalid_batch_request); + auto valid_slice = valid_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, 0, 1}); + ASSERT_NE(valid_slice, nullptr); + ASSERT_TRUE(valid_task->seal()); + auto unsealed_task = UbTask::create(invalid_batch_request); + ASSERT_NE(unsealed_task->addSlice(UbSliceSpec{ + source.data(), reinterpret_cast(target.data()), 16, + 0, 1}), + nullptr); + + EXPECT_FALSE(workers.submitBatch({valid_task, unsealed_task}).ok()); + EXPECT_EQ(workers.queuedCount(), 0U); + EXPECT_EQ(workers.inflightCount(), 0U); + EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kInitial); + adapter->holdNextCompletion(); adapter->dropNextQuiescedCompletion(); @@ -762,6 +787,15 @@ TEST(UbNativeDataPathTest, EndpointStoreNeverReusesRetiredGeneration) { EXPECT_EQ(concurrent[i], concurrent[0]); } auto first = concurrent[0]; + UbBootstrapDesc local_bootstrap; + ASSERT_TRUE(first + ->makeBootstrapDesc("local", "local@ub:fake0:eid0", + key.peer_nic_path, 1, local_bootstrap) + .ok()); + EXPECT_NE( + std::find(local_bootstrap.capabilities.begin(), + local_bootstrap.capabilities.end(), "receiver_credit_v1"), + local_bootstrap.capabilities.end()); const uint64_t old_generation = first->generation(); EXPECT_TRUE(store.retire(key, old_generation)); std::shared_ptr replacement; @@ -885,8 +919,13 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { ASSERT_TRUE(transport.addMemoryBuffer(descriptors, options).ok()); ASSERT_TRUE(control->segmentManager() .updateLocal([&](SegmentDesc& segment) { - std::get(segment.detail).buffers = - descriptors; + auto& memory = + std::get(segment.detail); + memory.buffers = descriptors; + ReceiverCreditAdvertV1 advert; + advert.receiver_session_id = {1, 2}; + advert.epoch = 1; + memory.receiver_credit = std::move(advert); return Status::OK(); }) .ok()); @@ -904,6 +943,13 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { transport.submitTransferTasks(rejected_batch, {oversized_request}) .IsInvalidArgument()); EXPECT_TRUE(transport.freeSubBatch(rejected_batch).ok()); + UbBootstrapDesc unsupported_bootstrap; + UbBootstrapDesc rejected_bootstrap; + ASSERT_TRUE(ControlClient::bootstrapUb(segment_name, unsupported_bootstrap, + rejected_bootstrap) + .ok()); + EXPECT_NE(rejected_bootstrap.reply_msg.find("receiver-credit"), + std::string::npos); Transport::SubBatchRef batch = nullptr; ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok()); @@ -913,6 +959,11 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) { request.target_id = LOCAL_SEGMENT_ID; request.target_offset = reinterpret_cast(target.data()); request.length = source.size(); + auto invalid_request = request; + invalid_request.length = 0; + EXPECT_FALSE( + transport.submitTransferTasks(batch, {request, invalid_request}).ok()); + EXPECT_EQ(batch->size(), 0U); ASSERT_TRUE(transport.submitTransferTasks(batch, {request}).ok()); TransferStatus transfer{PENDING, 0}; From c0ba30f3a85f4b664f441889799e663fbd98fb96 Mon Sep 17 00:00:00 2001 From: Connor-Matthew <60215777+Connor-Matthew@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:29:31 +0800 Subject: [PATCH 14/14] [TENT] Preserve receiver credit fencing on HP TCP after migration --- .../src/transport/hp_tcp/hp_tcp_transport.cpp | 19 +++++ .../tent/tests/hp_tcp_transport_test.cpp | 79 +++++++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp index 978bd7d934..335360c70f 100644 --- a/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/hp_tcp/hp_tcp_transport.cpp @@ -500,6 +500,25 @@ Status HighPerformanceTcpTransport::planTask( if (segment == nullptr || segment->type != SegmentType::Memory) { return NeedsRefresh("HP TCP target is not a memory segment"); } + const bool has_credit_fence = + batch.receiver_credit_session_high != 0 || + batch.receiver_credit_session_low != 0 || + batch.receiver_credit_epoch != 0; + if (has_credit_fence) { + const auto& advert = segment->getMemory().receiver_credit; + if (!advert || + advert->schema_version != kReceiverCreditProtocolVersion || + advert->flags != kReceiverCreditRequired || + advert->receiver_session_id.high != + batch.receiver_credit_session_high || + advert->receiver_session_id.low != + batch.receiver_credit_session_low || + advert->epoch != batch.receiver_credit_epoch) { + return Status::InvalidEntry( + "HP TCP receiver-credit session changed before " + "post" LOC_MARK); + } + } BufferDesc* buffer = segment->findBuffer(request.target_offset, request.length); if (buffer == nullptr || diff --git a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp index 452b2e3f04..abca485bc6 100644 --- a/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/hp_tcp_transport_test.cpp @@ -251,6 +251,85 @@ Status WaitForTransportResult(HighPerformanceTcpTransport& transport, return Status::InternalError("HP TCP test transfer did not finish"); } +TEST(HighPerformanceTcpTransportTest, ReceiverCreditFenceTracksPinnedMetadata) { + auto metadata = MakeLocalMetadata(); + HighPerformanceTcpTransport transport(MakeParams()); + std::string name = "hp_transport_test"; + ASSERT_TRUE(transport.install(name, metadata, nullptr, nullptr).ok()); + std::array source{}; + std::array target{}; + source.fill(0x5a); + BufferDesc source_desc{}; + source_desc.addr = reinterpret_cast(source.data()); + source_desc.length = source.size(); + source_desc.location = "cpu:0"; + BufferDesc target_desc{}; + target_desc.addr = reinterpret_cast(target.data()); + target_desc.length = target.size(); + target_desc.location = "cpu:0"; + MemoryOptions options; + options.type = HP_TCP; + options.perm = kGlobalReadWrite; + ASSERT_TRUE(transport.addMemoryBuffer(source_desc, options).ok()); + ASSERT_TRUE(transport.addMemoryBuffer(target_desc, options).ok()); + ASSERT_TRUE(PublishBuffers(metadata, {source_desc, target_desc}).ok()); + + Request request{}; + request.opcode = Request::WRITE; + request.source = source.data(); + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = target_desc.addr; + request.length = source.size(); + ReceiverCreditAdvertV1 advert; + advert.receiver_session_id = {1, 2}; + advert.epoch = 7; + // No advert, wrong session, wrong epoch, and unsupported protocol must + // reject before worker handoff; a matching pinned advert must still work. + for (int scenario = 0; scenario < 6; ++scenario) { + SCOPED_TRACE(scenario); + auto current = advert; + if (scenario == 1) ++current.receiver_session_id.low; + if (scenario == 2) ++current.epoch; + if (scenario == 3) ++current.schema_version; + if (scenario == 4) current.flags = 0; + ASSERT_TRUE(metadata->segmentManager() + .updateLocal([&](SegmentDesc& segment) -> Status { + auto& published = + segment.getMemory().receiver_credit; + if (scenario == 0) + published.reset(); + else + published = current; + return Status::OK(); + }) + .ok()); + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 1).ok()); + batch->receiver_credit_session_high = advert.receiver_session_id.high; + batch->receiver_credit_session_low = advert.receiver_session_id.low; + batch->receiver_credit_epoch = advert.epoch; + const Status submitted = + transport.submitTransferTasks(batch, {request}); + if (scenario < 5) { + EXPECT_TRUE(submitted.IsInvalidEntry()) << submitted.ToString(); + EXPECT_EQ(batch->size(), 0u); + EXPECT_TRUE(std::all_of(target.begin(), target.end(), + [](uint8_t byte) { return byte == 0; })); + } else { + ASSERT_TRUE(submitted.ok()) << submitted.ToString(); + TransferStatus status{}; + ASSERT_TRUE(WaitForTransportResult(transport, batch, status).ok()); + EXPECT_EQ(status.s, COMPLETED); + EXPECT_EQ(source, target); + } + EXPECT_TRUE(transport.freeSubBatch(batch).ok()); + } + EXPECT_TRUE(transport.quiesce().ok()); + EXPECT_TRUE(transport.removeMemoryBuffer(source_desc).ok()); + EXPECT_TRUE(transport.removeMemoryBuffer(target_desc).ok()); + EXPECT_TRUE(transport.uninstall().ok()); +} + // Test-only relay: observe the wire request and forward it to the real server. // Pumped by the test thread; no transport worker waits for another callback. class RailRelay {