diff --git a/.dockerignore b/.dockerignore index 3f691b5bc..e0d9d8be6 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,6 +1,7 @@ # Keep build outputs, VCS data, and caches out of the docker build context. build/ .git/ +.rocprofv3/ **/__pycache__/ **/*.egg-info/ **/*.pyc diff --git a/src/umbp/CMakeLists.txt b/src/umbp/CMakeLists.txt index 96edbcfa5..68a47c367 100644 --- a/src/umbp/CMakeLists.txt +++ b/src/umbp/CMakeLists.txt @@ -308,10 +308,7 @@ add_library( ${UMBP_GRPC_SRCS} ${UMBP_PEER_PROTO_SRCS} ${UMBP_PEER_GRPC_SRCS} - distributed/master/global_block_index.cpp - distributed/master/external_kv_block_index.cpp - distributed/master/client_registry.cpp - distributed/master/external_kv_hit_index.cpp + distributed/master/in_memory_master_metadata_store.cpp distributed/master/master_server.cpp distributed/master/master_client.cpp distributed/master/rpc_latency_timer.cpp @@ -383,4 +380,6 @@ endif() target_link_libraries(umbp_client PRIVATE umbp_common ${_PROTOBUF_LIBS} ${_GRPCPP_LIB}) -add_subdirectory(tests) +# UMBP unit tests live under tests/cpp/umbp (built when BUILD_TESTS=ON); they +# were previously duplicated here under src/umbp/tests with a second GoogleTest +# fetch. diff --git a/src/umbp/distributed/bin/client_main.cpp b/src/umbp/distributed/bin/client_main.cpp index 6d1d6cd0c..528a82e96 100644 --- a/src/umbp/distributed/bin/client_main.cpp +++ b/src/umbp/distributed/bin/client_main.cpp @@ -99,6 +99,11 @@ int main(int argc, char** argv) { return 1; } + // This demo drives MasterClient directly (it does not go through PoolClient, + // which is what normally honors config.auto_heartbeat on Init). Start the + // heartbeat thread so the master's reaper does not expire us. + client.StartHeartbeat(); + constexpr auto kOperationInterval = std::chrono::seconds(3); uint64_t iteration = 0; diff --git a/src/umbp/distributed/master/client_registry.cpp b/src/umbp/distributed/master/client_registry.cpp deleted file mode 100644 index 76fefc281..000000000 --- a/src/umbp/distributed/master/client_registry.cpp +++ /dev/null @@ -1,298 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/master/client_registry.h" - -#include - -#include "mori/utils/mori_log.hpp" -#include "umbp/distributed/master/external_kv_block_index.h" -#include "umbp/distributed/master/global_block_index.h" - -namespace mori::umbp { - -ClientRegistry::ClientRegistry(const ClientRegistryConfig& config) : config_(config) {} - -ClientRegistry::ClientRegistry(const ClientRegistryConfig& config, GlobalBlockIndex& index, - ExternalKvBlockIndex* external_kv_index) - : config_(config), index_(&index), external_kv_index_(external_kv_index) {} - -ClientRegistry::~ClientRegistry() { StopReaper(); } - -void ClientRegistry::SetBlockIndex(GlobalBlockIndex* index) { - std::unique_lock lock(mutex_); - index_ = index; -} - -void ClientRegistry::SetExternalKvBlockIndex(ExternalKvBlockIndex* index) { - std::unique_lock lock(mutex_); - external_kv_index_ = index; -} - -bool ClientRegistry::RegisterClient(const std::string& node_id, const std::string& node_address, - const std::map& tier_capacities, - const std::string& peer_address, - const std::vector& engine_desc_bytes, - const std::vector& tags) { - std::unique_lock lock(mutex_); - const auto now = std::chrono::steady_clock::now(); - - auto it = clients_.find(node_id); - if (it != clients_.end()) { - const bool is_expired = (now - it->second.last_heartbeat > ExpiryDuration()) || - (it->second.status == ClientStatus::EXPIRED); - if (it->second.status == ClientStatus::ALIVE && !is_expired) { - MORI_UMBP_WARN("[Registry] Rejecting re-registration for alive node: {}", node_id); - return false; - } - MORI_UMBP_INFO("[Registry] Re-registering expired node: {}", node_id); - } - - ClientRecord record; - record.node_id = node_id; - record.node_address = node_address; - record.status = ClientStatus::ALIVE; - record.last_heartbeat = now; - record.registered_at = now; - record.tier_capacities = tier_capacities; - record.peer_address = peer_address; - record.engine_desc_bytes = engine_desc_bytes; - record.last_applied_seq = 0; - record.tags = tags; - - clients_[node_id] = std::move(record); - - std::string tags_str; - for (const auto& t : tags) { - if (!tags_str.empty()) tags_str += ','; - tags_str += t; - } - MORI_UMBP_INFO("[Registry] Registered node: {} at {} (peer={}) tags=[{}]", node_id, node_address, - peer_address, tags_str); - return true; -} - -void ClientRegistry::UnregisterClient(const std::string& node_id) { - GlobalBlockIndex* idx = nullptr; - ExternalKvBlockIndex* external_idx = nullptr; - { - std::unique_lock lock(mutex_); - auto it = clients_.find(node_id); - if (it == clients_.end()) return; - idx = index_; - external_idx = external_kv_index_; - clients_.erase(it); - } - if (idx != nullptr) { - idx->RemoveByNode(node_id); - } - if (external_idx != nullptr) { - external_idx->UnregisterByNode(node_id); - } - MORI_UMBP_INFO("[Registry] Unregistered node: {}", node_id); -} - -ClientStatus ClientRegistry::Heartbeat(const std::string& node_id, - const std::map& tier_capacities, - const std::vector& bundles, bool is_full_sync, - uint64_t delta_seq_baseline, uint64_t* out_acked_seq, - bool* out_request_full_sync) { - if (out_acked_seq != nullptr) *out_acked_seq = 0; - if (out_request_full_sync != nullptr) *out_request_full_sync = false; - - GlobalBlockIndex* idx = nullptr; - std::vector bundles_to_apply; - std::vector full_sync_adds; - bool do_full_sync = false; - - { - std::unique_lock lock(mutex_); - auto it = clients_.find(node_id); - if (it == clients_.end()) { - MORI_UMBP_WARN("[Registry] Heartbeat from unknown node: {}", node_id); - return ClientStatus::UNKNOWN; - } - auto& record = it->second; - - record.last_heartbeat = std::chrono::steady_clock::now(); - record.status = ClientStatus::ALIVE; - record.tier_capacities = tier_capacities; - - idx = index_; - - if (is_full_sync) { - for (const auto& bundle : bundles) { - for (auto ev : bundle.events) { - if (ev.kind != KvEvent::Kind::ADD) continue; - full_sync_adds.push_back(std::move(ev)); - } - } - record.last_applied_seq = delta_seq_baseline; - if (out_acked_seq != nullptr) *out_acked_seq = record.last_applied_seq; - do_full_sync = true; - } else { - for (const auto& bundle : bundles) { - if (bundle.seq <= record.last_applied_seq) continue; - if (bundle.seq != record.last_applied_seq + 1) { - MORI_UMBP_WARN( - "[Registry] Heartbeat bundle seq gap from {}: got {}, expected {} — requesting " - "full sync", - node_id, bundle.seq, record.last_applied_seq + 1); - if (out_acked_seq != nullptr) *out_acked_seq = record.last_applied_seq; - if (out_request_full_sync != nullptr) *out_request_full_sync = true; - return ClientStatus::ALIVE; - } - bundles_to_apply.push_back(bundle); - record.last_applied_seq = bundle.seq; - } - if (out_acked_seq != nullptr) *out_acked_seq = record.last_applied_seq; - } - } - - if (idx != nullptr) { - if (do_full_sync) { - idx->ReplaceNodeLocations(node_id, full_sync_adds); - } else { - for (const auto& bundle : bundles_to_apply) { - if (!bundle.events.empty()) idx->ApplyEvents(node_id, bundle.events); - } - } - } - return ClientStatus::ALIVE; -} - -bool ClientRegistry::IsClientAlive(const std::string& node_id) const { - std::shared_lock lock(mutex_); - auto it = clients_.find(node_id); - return it != clients_.end() && it->second.status == ClientStatus::ALIVE; -} - -size_t ClientRegistry::ClientCount() const { - std::shared_lock lock(mutex_); - return clients_.size(); -} - -size_t ClientRegistry::AliveClientCount() const { - std::shared_lock lock(mutex_); - size_t count = 0; - for (const auto& [id, record] : clients_) { - if (record.status == ClientStatus::ALIVE) ++count; - } - return count; -} - -std::vector ClientRegistry::GetAliveClients() const { - std::shared_lock lock(mutex_); - std::vector result; - for (const auto& [id, record] : clients_) { - if (record.status == ClientStatus::ALIVE) result.push_back(record); - } - return result; -} - -AlivePeerView ClientRegistry::GetAlivePeerView() const { - std::shared_lock lock(mutex_); - AlivePeerView view; - view.reserve(clients_.size()); - for (const auto& [id, record] : clients_) { - if (record.status != ClientStatus::ALIVE) continue; - view.emplace(record.node_id, record.peer_address); - } - return view; -} - -std::vector ClientRegistry::GetClientTags(const std::string& node_id) const { - std::shared_lock lock(mutex_); - auto it = clients_.find(node_id); - if (it == clients_.end()) return {}; - return it->second.tags; -} - -void ClientRegistry::StartReaper() { - reaper_running_ = true; - reaper_thread_ = std::thread(&ClientRegistry::ReaperLoop, this); - MORI_UMBP_INFO("[Reaper] Started (interval={}s, expiry={}s)", config_.reaper_interval.count(), - ExpiryDuration().count()); -} - -void ClientRegistry::StopReaper() { - if (reaper_running_) { - reaper_running_ = false; - reaper_cv_.notify_one(); - if (reaper_thread_.joinable()) reaper_thread_.join(); - MORI_UMBP_INFO("[Reaper] Stopped"); - } -} - -void ClientRegistry::ReaperLoop() { - while (reaper_running_) { - { - std::unique_lock cv_lock(reaper_cv_mutex_); - reaper_cv_.wait_for(cv_lock, config_.reaper_interval, - [this] { return !reaper_running_.load(); }); - } - if (!reaper_running_) break; - ReapExpiredClients(); - } -} - -void ClientRegistry::ReapExpiredClients() { - const auto now = std::chrono::steady_clock::now(); - const auto expiry = ExpiryDuration(); - std::vector dead_nodes; - - { - std::unique_lock lock(mutex_); - auto it = clients_.begin(); - while (it != clients_.end()) { - if (now - it->second.last_heartbeat > expiry) { - MORI_UMBP_WARN("[Reaper] Reaping expired client: {}", it->first); - dead_nodes.push_back(it->first); - it = clients_.erase(it); - } else { - ++it; - } - } - } - - GlobalBlockIndex* idx = nullptr; - ExternalKvBlockIndex* external_idx = nullptr; - { - std::shared_lock lock(mutex_); - idx = index_; - external_idx = external_kv_index_; - } - - if (idx != nullptr) { - for (const auto& dead_id : dead_nodes) { - // Clear every index entry belonging to the dead node. Capacity - // numbers vanish with the ClientRecord above. - idx->RemoveByNode(dead_id); - } - } - if (external_idx != nullptr) { - for (const auto& dead_id : dead_nodes) { - external_idx->UnregisterByNode(dead_id); - } - } -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/master/eviction_manager.cpp b/src/umbp/distributed/master/eviction_manager.cpp index 5a89e926a..729fb410d 100644 --- a/src/umbp/distributed/master/eviction_manager.cpp +++ b/src/umbp/distributed/master/eviction_manager.cpp @@ -21,8 +21,9 @@ // SOFTWARE. #include "umbp/distributed/master/eviction_manager.h" +#include #include -#include +#include #include #include #include @@ -30,17 +31,16 @@ #include #include "mori/utils/mori_log.hpp" -#include "umbp/distributed/master/client_registry.h" #include "umbp/distributed/master/evict_strategy.h" -#include "umbp/distributed/master/global_block_index.h" +#include "umbp/distributed/master/master_metadata_store.h" +#include "umbp/distributed/types.h" namespace mori::umbp { -EvictionManager::EvictionManager(GlobalBlockIndex& index, ClientRegistry& registry, - const EvictionConfig& config, EvictKeyDispatcher* dispatcher, +EvictionManager::EvictionManager(IMasterMetadataStore& store, const EvictionConfig& config, + EvictKeyDispatcher* dispatcher, std::unique_ptr strategy) - : index_(index), - registry_(registry), + : store_(store), config_(config), dispatcher_(dispatcher), strategy_(strategy ? std::move(strategy) : std::make_unique()) {} @@ -80,11 +80,13 @@ void EvictionManager::EvictionLoop() { // This function picks victims and dispatches EvictKey to each peer via the // dispatcher; master state itself is left untouched here. void EvictionManager::RunOnce() { - auto clients = registry_.GetAliveClients(); + auto clients = store_.ListAliveClients(); - using NodeTierKey = GlobalBlockIndex::NodeTierKey; - std::set overloaded_node_tiers; - std::unordered_map> bytes_to_free; + // Per-(node, tier) byte budget down to the LOW watermark. The budget map's + // keys double as the set of overloaded buckets to enumerate; the byte values + // are handed to the strategy (never to the store) for the actual victim + // decision. + std::map bytes_to_free; for (const auto& client : clients) { for (const auto& [tier, cap] : client.tier_capacities) { @@ -99,31 +101,50 @@ void EvictionManager::RunOnce() { if (usage >= config_.high_watermark) { auto target_used = static_cast(static_cast(cap.total_bytes) * config_.low_watermark); - auto to_free = static_cast(used) - static_cast(target_used); - if (to_free > 0) { - overloaded_node_tiers.insert({client.node_id, tier}); - bytes_to_free[client.node_id][tier] += to_free; + if (used > target_used) { + bytes_to_free[{client.node_id, tier}] += used - target_used; } } } } - if (overloaded_node_tiers.empty()) return; - MORI_UMBP_INFO("[EvictionManager] {} overloaded node-tiers detected", - overloaded_node_tiers.size()); - - auto candidates = index_.FindEvictionCandidates(overloaded_node_tiers); - if (candidates.empty()) { + if (bytes_to_free.empty()) return; + MORI_UMBP_INFO("[EvictionManager] {} overloaded node-tiers detected", bytes_to_free.size()); + + // Ask the store for eviction-eligible candidates in the overloaded buckets. + // The store is policy-neutral: it returns rows ordered by the hint (LRU here, + // so an indexed backend can push the ordering down) but makes no eviction + // decision and never sees the byte budget. max_per_bucket=0 → no cap, so the + // strategy gets full candidate visibility. + std::vector buckets; + buckets.reserve(bytes_to_free.size()); + for (const auto& [ntk, _bytes] : bytes_to_free) buckets.push_back(ntk); + + auto candidates_by_bucket = + store_.EnumerateEvictionCandidates(buckets, EvictionOrder::kLeastRecentlyAccessed, + /*max_per_bucket=*/0, std::chrono::system_clock::now()); + if (candidates_by_bucket.empty()) { MORI_UMBP_DEBUG("[EvictionManager] No eviction candidates found"); return; } - // Policy ranks candidates and picks victims within budget, grouped by node so - // each peer gets a single EvictKey keys[]. - auto per_node_keys = strategy_->SelectVictims(std::move(candidates), std::move(bytes_to_free)); + // Flatten the candidates and translate the per-(node, tier) byte budget into + // the strategy's shape, then let the strategy pick victims (default LRU), + // grouped by node so each peer gets a single EvictKey keys[]. + std::vector candidates; + for (auto& [bucket, rows] : candidates_by_bucket) { + for (auto& c : rows) candidates.push_back(std::move(c)); + } + std::unordered_map> strategy_budget; + for (const auto& [ntk, bytes] : bytes_to_free) { + strategy_budget[ntk.node_id][ntk.tier] = static_cast(bytes); + } + + auto per_node_keys = strategy_->SelectVictims(std::move(candidates), std::move(strategy_budget)); size_t selected = 0; for (const auto& [node_id, keys] : per_node_keys) selected += keys.size(); + if (selected == 0) return; MORI_UMBP_INFO("[EvictionManager] Selected {} victims across {} nodes", selected, diff --git a/src/umbp/distributed/master/external_kv_block_index.cpp b/src/umbp/distributed/master/external_kv_block_index.cpp deleted file mode 100644 index f6e93ac33..000000000 --- a/src/umbp/distributed/master/external_kv_block_index.cpp +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/master/external_kv_block_index.h" - -#include -#include -#include -#include - -namespace mori::umbp { - -size_t ExternalKvBlockIndex::Register(const std::string& node_id, - const std::vector& hashes, TierType tier) { - std::unique_lock lock(mutex_); - size_t mutated = 0; - for (const auto& hash : hashes) { - auto [it, inserted] = entries_[hash][node_id].insert(tier); - (void)it; - if (inserted) ++mutated; - } - return mutated; -} - -size_t ExternalKvBlockIndex::Unregister(const std::string& node_id, - const std::vector& hashes, TierType tier) { - std::unique_lock lock(mutex_); - size_t mutated = 0; - for (const auto& hash : hashes) { - auto it = entries_.find(hash); - if (it == entries_.end()) continue; - - auto node_it = it->second.find(node_id); - if (node_it == it->second.end()) continue; - - mutated += node_it->second.erase(tier); - if (node_it->second.empty()) it->second.erase(node_it); - if (it->second.empty()) entries_.erase(it); - } - return mutated; -} - -size_t ExternalKvBlockIndex::UnregisterByNodeAtTier(const std::string& node_id, TierType tier) { - std::unique_lock lock(mutex_); - size_t mutated = 0; - auto it = entries_.begin(); - while (it != entries_.end()) { - auto node_it = it->second.find(node_id); - if (node_it != it->second.end()) { - mutated += node_it->second.erase(tier); - if (node_it->second.empty()) it->second.erase(node_it); - } - if (it->second.empty()) { - it = entries_.erase(it); - } else { - ++it; - } - } - return mutated; -} - -size_t ExternalKvBlockIndex::UnregisterByNode(const std::string& node_id) { - std::unique_lock lock(mutex_); - size_t mutated = 0; - auto it = entries_.begin(); - while (it != entries_.end()) { - auto node_it = it->second.find(node_id); - if (node_it != it->second.end()) { - mutated += node_it->second.size(); - it->second.erase(node_it); - } - if (it->second.empty()) { - it = entries_.erase(it); - } else { - ++it; - } - } - return mutated; -} - -std::vector ExternalKvBlockIndex::Match( - const std::vector& hashes) const { - std::shared_lock lock(mutex_); - - std::unordered_map>> acc; - for (const auto& hash : hashes) { - auto it = entries_.find(hash); - if (it == entries_.end()) continue; - for (const auto& [node_id, tiers] : it->second) { - auto& by_tier = acc[node_id]; - for (TierType tier : tiers) by_tier[tier].push_back(hash); - } - } - - std::vector result; - result.reserve(acc.size()); - for (auto& [node_id, by_tier] : acc) { - NodeMatch m; - m.node_id = std::move(node_id); - m.hashes_by_tier = std::move(by_tier); - result.push_back(std::move(m)); - } - return result; -} - -size_t ExternalKvBlockIndex::GetKvCount(const std::string& node_id) const { - std::shared_lock lock(mutex_); - size_t count = 0; - for (const auto& [hash, nodes] : entries_) { - (void)hash; - if (nodes.count(node_id)) ++count; - } - return count; -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/master/external_kv_hit_index.cpp b/src/umbp/distributed/master/external_kv_hit_index.cpp deleted file mode 100644 index 84d775c0e..000000000 --- a/src/umbp/distributed/master/external_kv_hit_index.cpp +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/master/external_kv_hit_index.h" - -#include -#include -#include - -namespace mori::umbp { - -size_t ExternalKvHitIndex::ShardIdx(std::string_view hash) { - return std::hash{}(hash) % kShards; -} - -void ExternalKvHitIndex::UpdateLastSeen(Entry* entry, uint64_t now_ns) { - uint64_t old = entry->last_seen_ns.load(std::memory_order_relaxed); - while (old < now_ns && !entry->last_seen_ns.compare_exchange_weak( - old, now_ns, std::memory_order_relaxed, std::memory_order_relaxed)) { - } -} - -void ExternalKvHitIndex::IncrementHits(const std::vector& unique_hashes, - uint64_t now_ns) { - for (const auto& hash : unique_hashes) { - auto& shard = shards_[ShardIdx(hash)]; - { - std::shared_lock lock(shard.mu); - auto it = shard.entries.find(hash); - if (it != shard.entries.end()) { - Entry* entry = it->second.get(); - entry->total.fetch_add(1, std::memory_order_relaxed); - UpdateLastSeen(entry, now_ns); - continue; - } - } - - std::unique_lock lock(shard.mu); - auto [it, inserted] = shard.entries.try_emplace(hash); - if (inserted) { - auto entry = std::make_unique(); - entry->total.store(1, std::memory_order_relaxed); - entry->last_seen_ns.store(now_ns, std::memory_order_relaxed); - it->second = std::move(entry); - } else { - Entry* entry = it->second.get(); - entry->total.fetch_add(1, std::memory_order_relaxed); - UpdateLastSeen(entry, now_ns); - } - } -} - -size_t ExternalKvHitIndex::Lookup(const std::vector& hashes, - std::vector>* out) const { - if (out == nullptr) return 0; - std::unordered_set seen; - seen.reserve(hashes.size()); - size_t filled = 0; - for (const auto& hash : hashes) { - if (!seen.insert(hash).second) continue; - const auto& shard = shards_[ShardIdx(hash)]; - std::shared_lock lock(shard.mu); - auto it = shard.entries.find(hash); - if (it == shard.entries.end()) continue; - out->push_back({hash, it->second->total.load(std::memory_order_relaxed)}); - ++filled; - } - return filled; -} - -size_t ExternalKvHitIndex::GarbageCollect(uint64_t cutoff_ns) { - size_t dropped = 0; - for (auto& shard : shards_) { - std::unique_lock lock(shard.mu); - auto it = shard.entries.begin(); - while (it != shard.entries.end()) { - const uint64_t last_seen = it->second->last_seen_ns.load(std::memory_order_relaxed); - if (last_seen < cutoff_ns) { - it = shard.entries.erase(it); - ++dropped; - } else { - ++it; - } - } - } - return dropped; -} - -size_t ExternalKvHitIndex::Size() const { - size_t size = 0; - for (const auto& shard : shards_) { - std::shared_lock lock(shard.mu); - size += shard.entries.size(); - } - return size; -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/master/global_block_index.cpp b/src/umbp/distributed/master/global_block_index.cpp deleted file mode 100644 index 247f8c963..000000000 --- a/src/umbp/distributed/master/global_block_index.cpp +++ /dev/null @@ -1,381 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include "umbp/distributed/master/global_block_index.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "mori/utils/mori_log.hpp" - -namespace mori::umbp { - -namespace { - -// Default shard count. Picked so a typical 128-key heartbeat fans out to a -// handful of events per shard while keeping per-shard maps cache-friendly. -constexpr size_t kDefaultIndexShards = 32; - -// Upper bound so a fat-fingered env value can't try to allocate a runaway -// number of shards (each shard is a shared_mutex + two maps) and OOM at start. -constexpr size_t kMaxIndexShards = 4096; - -size_t IndexShardCount() { - if (const char* e = std::getenv("UMBP_MASTER_INDEX_SHARDS")) { - char* end = nullptr; - const long v = std::strtol(e, &end, 10); - if (end != e && v >= 1) { - const size_t shards = static_cast(v); - if (shards > kMaxIndexShards) { - MORI_UMBP_WARN("[GlobalBlockIndex] clamping UMBP_MASTER_INDEX_SHARDS={} to max {}", shards, - kMaxIndexShards); - return kMaxIndexShards; - } - return shards; - } - MORI_UMBP_WARN("[GlobalBlockIndex] ignoring invalid UMBP_MASTER_INDEX_SHARDS='{}'", e); - } - return kDefaultIndexShards; -} - -// Locate (or insert) the location for (node_id, tier) within an entry's -// location list. Caller MUST hold the unique lock. Returns a pointer -// into entry.locations that's stable until the next mutation. -std::pair FindOrInsertLocation(BlockEntry& entry, const std::string& node_id, - TierType tier) { - for (auto& loc : entry.locations) { - if (loc.node_id == node_id && loc.tier == tier) return {&loc, false}; - } - entry.locations.push_back(Location{node_id, /*size=*/0, tier}); - return {&entry.locations.back(), true}; -} - -bool HasLocationForNode(const BlockEntry& entry, const std::string& node_id) { - return std::any_of(entry.locations.begin(), entry.locations.end(), - [&](const Location& loc) { return loc.node_id == node_id; }); -} - -using EntryMap = std::unordered_map; -using NodeToKeys = std::unordered_map>; - -// Drop every location owned by `node_id` within one shard, driven by the -// shard-local reverse index so the cost is O(node's keys in this shard) rather -// than a full O(shard entries) scan. Caller holds the shard's unique lock. -void RemoveNodeLocationsLocked(EntryMap& entries, NodeToKeys& node_to_keys, - const std::string& node_id) { - auto rev_it = node_to_keys.find(node_id); - if (rev_it == node_to_keys.end()) return; - auto keys = std::move(rev_it->second); - node_to_keys.erase(rev_it); - for (const auto& key : keys) { - auto eit = entries.find(key); - if (eit == entries.end()) continue; - auto& locs = eit->second.locations; - locs.erase(std::remove_if(locs.begin(), locs.end(), - [&](const Location& l) { return l.node_id == node_id; }), - locs.end()); - if (locs.empty()) entries.erase(eit); - } -} - -// Apply a single ADD/REMOVE event to one shard's maps (caller holds the -// shard's unique lock). Returns 1 if it mutated a location, else 0. -size_t ApplyAddOrRemoveLocked(EntryMap& entries, NodeToKeys& node_to_keys, - const std::string& node_id, const KvEvent& ev, - std::chrono::steady_clock::time_point now) { - if (ev.kind == KvEvent::Kind::ADD) { - auto& entry = entries[ev.key]; - if (entry.locations.empty()) { - entry.metrics.created_at = now; - entry.metrics.last_accessed_at = now; - entry.metrics.access_count = 0; - entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); - entry.atomic_access_count.store(0, std::memory_order_relaxed); - } - auto [loc, inserted] = FindOrInsertLocation(entry, node_id, ev.tier); - // Idempotent; must run on duplicate ADDs too. - node_to_keys[node_id].insert(ev.key); - if (!inserted) { - MORI_UMBP_WARN( - "[GlobalBlockIndex] duplicate ADD for key='{}' node={} tier={} old_size={} " - "new_size={}; keeping existing location", - ev.key, node_id, TierTypeName(ev.tier), loc->size, ev.size); - return 0; - } - loc->size = ev.size; - return 1; - } - // REMOVE - auto it = entries.find(ev.key); - if (it == entries.end()) return 0; - auto& locs = it->second.locations; - const size_t before = locs.size(); - locs.erase( - std::remove_if(locs.begin(), locs.end(), - [&](const Location& l) { return l.node_id == node_id && l.tier == ev.tier; }), - locs.end()); - if (locs.size() == before) return 0; - // find(), not operator[]: don't grow an empty bucket for strangers. - if (!HasLocationForNode(it->second, node_id)) { - auto rev_it = node_to_keys.find(node_id); - if (rev_it != node_to_keys.end()) { - rev_it->second.erase(ev.key); - if (rev_it->second.empty()) node_to_keys.erase(rev_it); - } - } - if (locs.empty()) entries.erase(it); - return 1; -} - -} // namespace - -GlobalBlockIndex::GlobalBlockIndex() : num_shards_(IndexShardCount()) { - shards_.reserve(num_shards_); - for (size_t i = 0; i < num_shards_; ++i) shards_.push_back(std::make_unique()); -} - -size_t GlobalBlockIndex::ApplyEvents(const std::string& node_id, - const std::vector& events) { - if (events.empty()) return 0; - const auto now = std::chrono::steady_clock::now(); - - // Sort one (shard, original index) array into per-shard runs; the index - // tie-break keeps same-key events in order. Each shard is locked once for - // its contiguous run, leaving other shards readable concurrently. - std::vector> shard_order(events.size()); - for (size_t i = 0; i < events.size(); ++i) shard_order[i] = {shard_index(events[i].key), i}; - std::sort(shard_order.begin(), shard_order.end()); - - size_t mutated = 0; - for (size_t run_begin = 0; run_begin < shard_order.size();) { - const size_t si = shard_order[run_begin].first; - auto& sh = shard_at(si); - std::unique_lock lock(sh.mutex); - size_t run_end = run_begin; - for (; run_end < shard_order.size() && shard_order[run_end].first == si; ++run_end) { - const KvEvent& ev = events[shard_order[run_end].second]; - mutated += ApplyAddOrRemoveLocked(sh.entries, sh.node_to_keys, node_id, ev, now); - } - run_begin = run_end; - } - return mutated; -} - -void GlobalBlockIndex::ReplaceNodeLocations(const std::string& node_id, - const std::vector& adds) { - const auto now = std::chrono::steady_clock::now(); - - // Group ADDs by destination shard up front so each shard is locked once. - std::vector> adds_by_shard(num_shards_); - for (const auto& ev : adds) { - if (ev.kind != KvEvent::Kind::ADD) continue; - adds_by_shard[shard_index(ev.key)].push_back(&ev); - } - - // Each shard owns exactly the node's keys that hash to it, so clearing the - // node + replaying its ADDs can be done shard-local with one lock per shard. - // This replaces the old single giant exclusive critical section with - // num_shards_ small ones — the win for full-sync tail latency. - for (size_t si = 0; si < num_shards_; ++si) { - auto& sh = shard_at(si); - std::unique_lock lock(sh.mutex); - - // O(N_node_in_shard) clear via the shard-local reverse index. - RemoveNodeLocationsLocked(sh.entries, sh.node_to_keys, node_id); - - for (const KvEvent* ev : adds_by_shard[si]) { - auto& entry = sh.entries[ev->key]; - if (entry.locations.empty()) { - entry.metrics.created_at = now; - entry.metrics.last_accessed_at = now; - entry.metrics.access_count = 0; - entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); - entry.atomic_access_count.store(0, std::memory_order_relaxed); - } - auto [loc, inserted] = FindOrInsertLocation(entry, node_id, ev->tier); - (void)inserted; - loc->size = ev->size; - sh.node_to_keys[node_id].insert(ev->key); - } - } -} - -void GlobalBlockIndex::RemoveByNode(const std::string& node_id) { - // O(N_node_in_shard) per shard via the reverse index (not a full scan). - for (size_t si = 0; si < num_shards_; ++si) { - auto& sh = shard_at(si); - std::unique_lock lock(sh.mutex); - RemoveNodeLocationsLocked(sh.entries, sh.node_to_keys, node_id); - } -} - -void GlobalBlockIndex::RecordAccess(const std::string& key) { - auto& sh = shard_for(key); - std::shared_lock lock(sh.mutex); - auto it = sh.entries.find(key); - if (it == sh.entries.end()) return; - it->second.RecordAccessAtomic(); -} - -void GlobalBlockIndex::GrantLease(const std::string& key, - std::chrono::steady_clock::duration duration) { - auto& sh = shard_for(key); - std::shared_lock lock(sh.mutex); - auto it = sh.entries.find(key); - if (it != sh.entries.end()) it->second.GrantLease(duration); -} - -std::vector GlobalBlockIndex::Lookup(const std::string& key) const { - auto& sh = shard_for(key); - std::shared_lock lock(sh.mutex); - auto it = sh.entries.find(key); - if (it == sh.entries.end()) return {}; - return it->second.locations; -} - -std::vector GlobalBlockIndex::BatchLookupExists(const std::vector& keys) const { - std::vector results(keys.size(), false); - if (keys.empty()) return results; - // Single-key fast path (RoutePut/RouteGet route through size-1 batches): - // skip the per-shard grouping vectors and lock just the one shard. - if (keys.size() == 1) { - auto& sh = shard_for(keys[0]); - std::shared_lock lock(sh.mutex); - auto it = sh.entries.find(keys[0]); - results[0] = (it != sh.entries.end()) && !it->second.locations.empty(); - return results; - } - // Group key indices by shard so each shard's shared lock is taken once. - std::vector> idx_by_shard(num_shards_); - for (size_t i = 0; i < keys.size(); ++i) idx_by_shard[shard_index(keys[i])].push_back(i); - for (size_t si = 0; si < num_shards_; ++si) { - if (idx_by_shard[si].empty()) continue; - auto& sh = shard_at(si); - std::shared_lock lock(sh.mutex); - for (size_t i : idx_by_shard[si]) { - auto it = sh.entries.find(keys[i]); - results[i] = (it != sh.entries.end()) && !it->second.locations.empty(); - } - } - return results; -} - -std::optional GlobalBlockIndex::GetMetrics(const std::string& key) const { - auto& sh = shard_for(key); - std::shared_lock lock(sh.mutex); - auto it = sh.entries.find(key); - if (it == sh.entries.end()) return std::nullopt; - BlockMetrics result = it->second.metrics; - result.last_accessed_at = it->second.GetLastAccessed(); - result.access_count = it->second.atomic_access_count.load(std::memory_order_acquire); - return result; -} - -std::vector> GlobalBlockIndex::BatchLookupForRouteGet( - const std::vector& keys, const std::unordered_set& exclude_nodes, - std::chrono::steady_clock::duration lease_duration) { - std::vector> out(keys.size()); - if (keys.empty()) return out; - const bool has_exclude = !exclude_nodes.empty(); - // Read the clock once for the whole batch: the per-hit access/lease bump - // would otherwise call steady_clock::now() twice per key. All keys in a - // batch are stamped within the same ~ms, which is well within lease/LRU - // granularity. - const auto now = std::chrono::steady_clock::now(); - - // Per-key body, shared across the single-key fast path and the grouped path. - auto resolve_one = [&](Shard& sh, size_t i) { - auto it = sh.entries.find(keys[i]); - if (it == sh.entries.end()) return; - const auto& src = it->second.locations; - auto& locs = out[i]; - if (!has_exclude) { - // Common case: copy the whole location vector in one allocation - // (no per-element vector growth). - locs = src; - } else { - locs.reserve(src.size()); - for (const auto& loc : src) { - if (exclude_nodes.count(loc.node_id)) continue; - locs.push_back(loc); - } - } - if (locs.empty()) return; - it->second.RecordAccessAtomic(now); - it->second.GrantLease(now, lease_duration); - }; - - // Single-key fast path (RoutePut/RouteGet route through size-1 batches): - // skip the per-shard grouping vectors and lock just the one shard. - if (keys.size() == 1) { - auto& sh = shard_for(keys[0]); - std::shared_lock lock(sh.mutex); - resolve_one(sh, 0); - return out; - } - - std::vector> idx_by_shard(num_shards_); - for (size_t i = 0; i < keys.size(); ++i) idx_by_shard[shard_index(keys[i])].push_back(i); - for (size_t si = 0; si < num_shards_; ++si) { - if (idx_by_shard[si].empty()) continue; - auto& sh = shard_at(si); - std::shared_lock lock(sh.mutex); - for (size_t i : idx_by_shard[si]) resolve_one(sh, i); - } - return out; -} - -std::vector GlobalBlockIndex::FindEvictionCandidates( - const std::set& overloaded_node_tiers) const { - // Best-effort advisory scan: lock one shard at a time (not a global atomic - // snapshot). Candidates are re-validated downstream, so a concurrent apply - // on an already-scanned shard is harmless. - std::vector candidates; - for (size_t si = 0; si < num_shards_; ++si) { - auto& sh = shard_at(si); - std::shared_lock lock(sh.mutex); - for (const auto& [key, entry] : sh.entries) { - if (entry.IsLeased()) continue; - for (const auto& loc : entry.locations) { - if (overloaded_node_tiers.count({loc.node_id, loc.tier})) { - EvictionCandidate c; - c.key = key; - c.location = loc; - c.last_accessed_at = entry.GetLastAccessed(); - c.size = loc.size; - candidates.push_back(std::move(c)); - } - } - } - } - return candidates; -} - -} // namespace mori::umbp diff --git a/src/umbp/distributed/master/in_memory_master_metadata_store.cpp b/src/umbp/distributed/master/in_memory_master_metadata_store.cpp new file mode 100644 index 000000000..08c5c24c1 --- /dev/null +++ b/src/umbp/distributed/master/in_memory_master_metadata_store.cpp @@ -0,0 +1,774 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#include "umbp/distributed/master/in_memory_master_metadata_store.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mori/utils/mori_log.hpp" + +namespace mori::umbp { + +namespace { + +// Default shard count. Picked so a typical 128-key heartbeat fans out to a +// handful of events per shard while keeping per-shard maps cache-friendly. +constexpr size_t kDefaultIndexShards = 32; + +// Upper bound so a fat-fingered env value can't try to allocate a runaway +// number of shards (each shard is a shared_mutex + two maps) and OOM at start. +constexpr size_t kMaxIndexShards = 4096; + +size_t IndexShardCount() { + if (const char* e = std::getenv("UMBP_MASTER_INDEX_SHARDS")) { + char* end = nullptr; + const long v = std::strtol(e, &end, 10); + if (end != e && v >= 1) { + const size_t shards = static_cast(v); + if (shards > kMaxIndexShards) { + MORI_UMBP_WARN("[MetadataStore] clamping UMBP_MASTER_INDEX_SHARDS={} to max {}", shards, + kMaxIndexShards); + return kMaxIndexShards; + } + return shards; + } + MORI_UMBP_WARN("[MetadataStore] ignoring invalid UMBP_MASTER_INDEX_SHARDS='{}'", e); + } + return kDefaultIndexShards; +} + +// Locate (or insert) the location for (node_id, tier) within a location list. +// Caller MUST hold the shard's unique lock. Returns a pointer into `locations` +// that's stable until the next mutation. +std::pair FindOrInsertLocation(std::vector& locations, + const std::string& node_id, TierType tier) { + for (auto& loc : locations) { + if (loc.node_id == node_id && loc.tier == tier) return {&loc, false}; + } + locations.push_back(Location{node_id, /*size=*/0, tier}); + return {&locations.back(), true}; +} + +bool HasLocationForNode(const std::vector& locations, const std::string& node_id) { + return std::any_of(locations.begin(), locations.end(), + [&](const Location& loc) { return loc.node_id == node_id; }); +} + +} // namespace + +InMemoryMasterMetadataStore::InMemoryMasterMetadataStore() : num_shards_(IndexShardCount()) { + shards_.reserve(num_shards_); + for (size_t i = 0; i < num_shards_; ++i) shards_.push_back(std::make_unique()); +} + +// ===================================================================== +// Single-shard block helpers (caller holds the shard's unique lock) +// ===================================================================== + +size_t InMemoryMasterMetadataStore::ApplyAddOrRemoveLocked( + Shard& shard, const std::string& node_id, const KvEvent& ev, + std::chrono::system_clock::time_point now) { + if (ev.kind == KvEvent::Kind::ADD) { + auto& entry = shard.entries[ev.key]; + if (entry.locations.empty()) { + entry.metrics.created_at = now; + entry.metrics.last_accessed_at = now; + entry.metrics.access_count = 0; + entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); + entry.atomic_access_count.store(0, std::memory_order_relaxed); + } + auto [loc, inserted] = FindOrInsertLocation(entry.locations, node_id, ev.tier); + // Idempotent; must run on duplicate ADDs too. + shard.node_to_keys[node_id].insert(ev.key); + if (!inserted) { + MORI_UMBP_WARN( + "[MetadataStore] duplicate ADD for key='{}' node={} tier={} old_size={} " + "new_size={}; keeping existing location", + ev.key, node_id, TierTypeName(ev.tier), loc->size, ev.size); + return 0; + } + loc->size = ev.size; + return 1; + } + // REMOVE + auto it = shard.entries.find(ev.key); + if (it == shard.entries.end()) return 0; + auto& locs = it->second.locations; + const size_t before = locs.size(); + locs.erase( + std::remove_if(locs.begin(), locs.end(), + [&](const Location& l) { return l.node_id == node_id && l.tier == ev.tier; }), + locs.end()); + if (locs.size() == before) return 0; + // find(), not operator[]: don't grow an empty bucket for strangers. + if (!HasLocationForNode(it->second.locations, node_id)) { + auto rev_it = shard.node_to_keys.find(node_id); + if (rev_it != shard.node_to_keys.end()) { + rev_it->second.erase(ev.key); + if (rev_it->second.empty()) shard.node_to_keys.erase(rev_it); + } + } + if (locs.empty()) shard.entries.erase(it); + return 1; +} + +void InMemoryMasterMetadataStore::RemoveNodeLocationsLocked(Shard& shard, + const std::string& node_id) { + auto rev_it = shard.node_to_keys.find(node_id); + if (rev_it == shard.node_to_keys.end()) return; + auto keys = std::move(rev_it->second); + shard.node_to_keys.erase(rev_it); + for (const auto& key : keys) { + auto eit = shard.entries.find(key); + if (eit == shard.entries.end()) continue; + auto& locs = eit->second.locations; + locs.erase(std::remove_if(locs.begin(), locs.end(), + [&](const Location& l) { return l.node_id == node_id; }), + locs.end()); + if (locs.empty()) shard.entries.erase(eit); + } +} + +// ===================================================================== +// Block-index writes (acquire shard lock(s) internally; NO meta_mutex_ held) +// ===================================================================== + +size_t InMemoryMasterMetadataStore::ApplyEventsToShards(const std::string& node_id, + const std::vector& events, + std::chrono::system_clock::time_point now) { + if (events.empty()) return 0; + + // Sort one (shard, original index) array into per-shard runs; the index + // tie-break keeps same-key events in order. Each shard is locked once for its + // contiguous run, leaving other shards readable concurrently. + std::vector> shard_order(events.size()); + for (size_t i = 0; i < events.size(); ++i) shard_order[i] = {shard_index(events[i].key), i}; + std::sort(shard_order.begin(), shard_order.end()); + + size_t mutated = 0; + for (size_t run_begin = 0; run_begin < shard_order.size();) { + const size_t si = shard_order[run_begin].first; + auto& sh = shard_at(si); + std::unique_lock lock(sh.mutex); + size_t run_end = run_begin; + for (; run_end < shard_order.size() && shard_order[run_end].first == si; ++run_end) { + const KvEvent& ev = events[shard_order[run_end].second]; + mutated += ApplyAddOrRemoveLocked(sh, node_id, ev, now); + } + run_begin = run_end; + } + return mutated; +} + +void InMemoryMasterMetadataStore::ReplaceNodeLocationsInShards( + const std::string& node_id, const std::vector& adds, + std::chrono::system_clock::time_point now) { + // Group ADDs by destination shard up front so each shard is locked once. + std::vector> adds_by_shard(num_shards_); + for (const auto& ev : adds) { + if (ev.kind != KvEvent::Kind::ADD) continue; + adds_by_shard[shard_index(ev.key)].push_back(&ev); + } + + // Each shard owns exactly the node's keys that hash to it, so clearing the + // node + replaying its ADDs is shard-local with one lock per shard. This is + // NOT a global snapshot: full_sync is atomic within each shard (a single key + // never tears), but a concurrent reader may observe some shards replaced and + // others not — see the atomicity contract in master_metadata_store.h. + for (size_t si = 0; si < num_shards_; ++si) { + auto& sh = shard_at(si); + std::unique_lock lock(sh.mutex); + + // O(N_node_in_shard) clear via the shard-local reverse index. + RemoveNodeLocationsLocked(sh, node_id); + + for (const KvEvent* ev : adds_by_shard[si]) { + auto& entry = sh.entries[ev->key]; + if (entry.locations.empty()) { + entry.metrics.created_at = now; + entry.metrics.last_accessed_at = now; + entry.metrics.access_count = 0; + entry.last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); + entry.atomic_access_count.store(0, std::memory_order_relaxed); + } + auto [loc, inserted] = FindOrInsertLocation(entry.locations, node_id, ev->tier); + (void)inserted; + loc->size = ev->size; + sh.node_to_keys[node_id].insert(ev->key); + } + } +} + +void InMemoryMasterMetadataStore::RemoveBlocksByNodeInShards(const std::string& node_id) { + // O(N_node_in_shard) per shard via the reverse index (not a full scan). + // UnregisterClient/ExpireStaleClients hit this per dead node. + for (size_t si = 0; si < num_shards_; ++si) { + auto& sh = shard_at(si); + std::unique_lock lock(sh.mutex); + RemoveNodeLocationsLocked(sh, node_id); + } +} + +// ===================================================================== +// Meta-domain locked helpers +// ===================================================================== + +void InMemoryMasterMetadataStore::RemoveExternalKvByNodeLocked(const std::string& node_id) { + auto it = external_kv_entries_.begin(); + while (it != external_kv_entries_.end()) { + it->second.erase(node_id); + if (it->second.empty()) { + it = external_kv_entries_.erase(it); + } else { + ++it; + } + } +} + +bool InMemoryMasterMetadataStore::IsClientAliveLocked(const std::string& node_id) const { + auto it = clients_.find(node_id); + return it != clients_.end() && it->second.status == ClientStatus::ALIVE; +} + +// ===================================================================== +// Cross-store writes +// ===================================================================== + +bool InMemoryMasterMetadataStore::RegisterClient(const ClientRegistration& registration, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after) { + std::unique_lock lock(meta_mutex_); + + auto it = clients_.find(registration.node_id); + if (it != clients_.end()) { + const bool is_stale = (now - it->second.last_heartbeat > stale_after) || + (it->second.status == ClientStatus::EXPIRED); + if (it->second.status == ClientStatus::ALIVE && !is_stale) { + MORI_UMBP_WARN("[MetadataStore] Rejecting re-registration for alive node: {}", + registration.node_id); + return false; + } + MORI_UMBP_INFO("[MetadataStore] Re-registering stale/expired node: {}", registration.node_id); + } + + ClientRecord record; + record.node_id = registration.node_id; + record.node_address = registration.node_address; + record.status = ClientStatus::ALIVE; + record.last_heartbeat = now; + record.registered_at = now; + record.tier_capacities = registration.tier_capacities; + record.peer_address = registration.peer_address; + record.engine_desc_bytes = registration.engine_desc_bytes; + record.last_applied_seq = 0; + record.tags = registration.tags; + + clients_[registration.node_id] = std::move(record); + + std::string tags_str; + for (const auto& t : registration.tags) { + if (!tags_str.empty()) tags_str += ','; + tags_str += t; + } + MORI_UMBP_INFO("[MetadataStore] Registered node: {} at {} (peer={}) tags=[{}]", + registration.node_id, registration.node_address, registration.peer_address, + tags_str); + return true; +} + +void InMemoryMasterMetadataStore::UnregisterClient(const std::string& node_id) { + // Cross-domain write (approach A): meta section (erase client + drop + // external-kv) under meta_mutex_, then the block wipe AFTER releasing it — + // never nested. Not globally atomic; the brief residue is benign (see the + // atomicity contract in master_metadata_store.h). + { + std::unique_lock lock(meta_mutex_); + auto it = clients_.find(node_id); + if (it == clients_.end()) return; // unknown node: nothing to cascade + clients_.erase(it); + RemoveExternalKvByNodeLocked(node_id); + } + // meta_mutex_ released (client existed, else we returned early): wipe its + // block locations per-shard. + RemoveBlocksByNodeInShards(node_id); + MORI_UMBP_INFO("[MetadataStore] Unregistered node: {}", node_id); +} + +HeartbeatResult InMemoryMasterMetadataStore::ApplyHeartbeat( + const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const std::map& caps, const std::vector& events, + bool is_full_sync) { + // Cross-domain write (approach A), hot path: seq-CAS + client-record update + // under meta_mutex_, then the block apply AFTER releasing it so heartbeats to + // different shards apply in parallel (PR #440's win) instead of serializing on + // the meta lock. Never nested; not globally atomic. See the atomicity contract + // in master_metadata_store.h. + bool apply = false; + { + std::unique_lock lock(meta_mutex_); + + auto it = clients_.find(node_id); + if (it == clients_.end()) { + MORI_UMBP_WARN("[MetadataStore] Heartbeat from unknown node: {}", node_id); + return HeartbeatResult{HeartbeatResult::UNKNOWN, 0}; + } + auto& record = it->second; + + // Gap check (CAS) on the delta path only — full_sync replaces wholesale and + // re-baselines last_applied_seq. + if (!is_full_sync && seq != record.last_applied_seq + 1) { + // SEQ_GAP: keep the node alive (it IS heartbeating, just mid-recovery) but + // do NOT advance caps or last_applied_seq. See hazard #1. + MORI_UMBP_WARN( + "[MetadataStore] Heartbeat seq gap from {}: got {}, expected {} — requesting full sync", + node_id, seq, record.last_applied_seq + 1); + record.last_heartbeat = now; + record.status = ClientStatus::ALIVE; + return HeartbeatResult{HeartbeatResult::SEQ_GAP, record.last_applied_seq}; + } + + record.last_heartbeat = now; + record.status = ClientStatus::ALIVE; + record.tier_capacities = caps; + record.last_applied_seq = seq; + apply = true; + } + + // meta_mutex_ released: apply events to the block shards without it held. + if (apply) { + if (is_full_sync) { + ReplaceNodeLocationsInShards(node_id, events, now); + } else { + ApplyEventsToShards(node_id, events, now); + } + } + return HeartbeatResult{HeartbeatResult::APPLIED, seq}; +} + +std::vector InMemoryMasterMetadataStore::ExpireStaleClients( + std::chrono::system_clock::time_point cutoff) { + // Cross-domain write (approach A): flip ALIVE→EXPIRED + drop external-kv under + // meta_mutex_; drop block locations AFTER releasing it, per-shard. Non-nested, + // not globally atomic (same benign residue window as UnregisterClient). + std::vector dead_nodes; + { + std::unique_lock lock(meta_mutex_); + for (auto& [node_id, record] : clients_) { + // Only ALIVE rows can transition to EXPIRED; an already-EXPIRED row is + // left alone so re-ticking the reaper is idempotent (its locations are + // already gone). EXPIRED rows are KEPT, not erased — see hazard #3. + if (record.status == ClientStatus::ALIVE && record.last_heartbeat < cutoff) { + MORI_UMBP_WARN("[MetadataStore] Expiring stale client: {}", node_id); + record.status = ClientStatus::EXPIRED; + dead_nodes.push_back(node_id); + } + } + for (const auto& dead_id : dead_nodes) RemoveExternalKvByNodeLocked(dead_id); + } + + // meta_mutex_ released: clear each dead node's block locations per-shard. + for (const auto& dead_id : dead_nodes) RemoveBlocksByNodeInShards(dead_id); + return dead_nodes; +} + +// ===================================================================== +// External-KV writes +// ===================================================================== + +bool InMemoryMasterMetadataStore::RegisterExternalKvIfAlive(const std::string& node_id, + const std::vector& hashes, + TierType tier) { + std::unique_lock lock(meta_mutex_); + if (!IsClientAliveLocked(node_id)) return false; + for (const auto& hash : hashes) { + external_kv_entries_[hash][node_id].insert(tier); + } + return true; +} + +void InMemoryMasterMetadataStore::UnregisterExternalKv(const std::string& node_id, + const std::vector& hashes, + TierType tier) { + std::unique_lock lock(meta_mutex_); + for (const auto& hash : hashes) { + auto it = external_kv_entries_.find(hash); + if (it == external_kv_entries_.end()) continue; + auto node_it = it->second.find(node_id); + if (node_it == it->second.end()) continue; + node_it->second.erase(tier); + if (node_it->second.empty()) it->second.erase(node_it); + if (it->second.empty()) external_kv_entries_.erase(it); + } +} + +void InMemoryMasterMetadataStore::UnregisterExternalKvByTier(const std::string& node_id, + TierType tier) { + std::unique_lock lock(meta_mutex_); + auto it = external_kv_entries_.begin(); + while (it != external_kv_entries_.end()) { + auto node_it = it->second.find(node_id); + if (node_it != it->second.end()) { + node_it->second.erase(tier); + if (node_it->second.empty()) it->second.erase(node_it); + } + if (it->second.empty()) { + it = external_kv_entries_.erase(it); + } else { + ++it; + } + } +} + +void InMemoryMasterMetadataStore::UnregisterExternalKvByNode(const std::string& node_id) { + std::unique_lock lock(meta_mutex_); + RemoveExternalKvByNodeLocked(node_id); +} + +std::size_t InMemoryMasterMetadataStore::GarbageCollectHits( + std::chrono::system_clock::time_point cutoff) { + std::unique_lock lock(meta_mutex_); + std::size_t dropped = 0; + auto it = external_kv_hits_.begin(); + while (it != external_kv_hits_.end()) { + if (it->second.last_seen < cutoff) { + it = external_kv_hits_.erase(it); + ++dropped; + } else { + ++it; + } + } + return dropped; +} + +// ===================================================================== +// Block reads +// ===================================================================== + +std::vector InMemoryMasterMetadataStore::LookupBlock(const std::string& key) const { + auto& sh = shard_for(key); + std::shared_lock lock(sh.mutex); + auto it = sh.entries.find(key); + if (it == sh.entries.end()) return {}; + return it->second.locations; +} + +std::vector InMemoryMasterMetadataStore::LookupBlockForRouteGet( + const std::string& key, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, std::chrono::system_clock::duration lease_duration) { + auto& sh = shard_for(key); + std::shared_lock lock(sh.mutex); + auto it = sh.entries.find(key); + if (it == sh.entries.end()) return {}; + + std::vector out; + for (const auto& loc : it->second.locations) { + if (!exclude_nodes.empty() && exclude_nodes.count(loc.node_id)) continue; + out.push_back(loc); + } + if (out.empty()) return out; + it->second.RecordAccessAtomic(now); + it->second.GrantLease(now, lease_duration); + return out; +} + +std::vector> InMemoryMasterMetadataStore::BatchLookupBlockForRouteGet( + const std::vector& keys, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, std::chrono::system_clock::duration lease_duration) { + std::vector> out(keys.size()); + if (keys.empty()) return out; + + // Per-key body, shared across the single-key fast path and the grouped path. + auto resolve_one = [&](Shard& sh, size_t i) { + auto it = sh.entries.find(keys[i]); + if (it == sh.entries.end()) return; + auto& locs = out[i]; + for (const auto& loc : it->second.locations) { + if (!exclude_nodes.empty() && exclude_nodes.count(loc.node_id)) continue; + locs.push_back(loc); + } + if (locs.empty()) return; + it->second.RecordAccessAtomic(now); + it->second.GrantLease(now, lease_duration); + }; + + // Single-key fast path (RoutePut/RouteGet route through size-1 batches): skip + // the per-shard grouping vectors and lock just the one shard. + if (keys.size() == 1) { + auto& sh = shard_for(keys[0]); + std::shared_lock lock(sh.mutex); + resolve_one(sh, 0); + return out; + } + + // Group key indices by shard so each shard's shared lock is taken once. + std::vector> idx_by_shard(num_shards_); + for (size_t i = 0; i < keys.size(); ++i) idx_by_shard[shard_index(keys[i])].push_back(i); + for (size_t si = 0; si < num_shards_; ++si) { + if (idx_by_shard[si].empty()) continue; + auto& sh = shard_at(si); + std::shared_lock lock(sh.mutex); + for (size_t i : idx_by_shard[si]) resolve_one(sh, i); + } + return out; +} + +std::vector InMemoryMasterMetadataStore::BatchExistsBlock( + const std::vector& keys) const { + std::vector results(keys.size(), false); + if (keys.empty()) return results; + + // Single-key fast path (RoutePut/RouteGet route through size-1 batches): skip + // the per-shard grouping vectors and lock just the one shard. + if (keys.size() == 1) { + auto& sh = shard_for(keys[0]); + std::shared_lock lock(sh.mutex); + auto it = sh.entries.find(keys[0]); + results[0] = (it != sh.entries.end()) && !it->second.locations.empty(); + return results; + } + + // Group key indices by shard so each shard's shared lock is taken once. + std::vector> idx_by_shard(num_shards_); + for (size_t i = 0; i < keys.size(); ++i) idx_by_shard[shard_index(keys[i])].push_back(i); + for (size_t si = 0; si < num_shards_; ++si) { + if (idx_by_shard[si].empty()) continue; + auto& sh = shard_at(si); + std::shared_lock lock(sh.mutex); + for (size_t i : idx_by_shard[si]) { + auto it = sh.entries.find(keys[i]); + results[i] = (it != sh.entries.end()) && !it->second.locations.empty(); + } + } + return results; +} + +std::map> +InMemoryMasterMetadataStore::EnumerateEvictionCandidates( + const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now) const { + std::map> result; + if (buckets.empty()) return result; + + // Requested buckets as a set for O(log n) membership during the scan. + const std::set wanted(buckets.begin(), buckets.end()); + + // 1. Best-effort advisory scan: lock ONE shard at a time (not a global atomic + // snapshot). Candidates are re-validated downstream, so a concurrent apply + // on an already-scanned shard is harmless. No maintained LRU index — each + // shard scan reads its entries directly, so tie-timestamp candidates are + // never dropped (§2d, Option A). Policy-neutral: every eligible row is + // collected; the byte budget and victim choice live in the strategy. All + // shard locks are released before the sort/cap step below (hard + // constraint). + for (size_t si = 0; si < num_shards_; ++si) { + auto& sh = shard_at(si); + std::shared_lock lock(sh.mutex); + for (const auto& [key, entry] : sh.entries) { + if (entry.IsLeased(now)) continue; + const auto last_accessed = entry.GetLastAccessed(); + for (const auto& loc : entry.locations) { + NodeTierKey ntk{loc.node_id, loc.tier}; + if (wanted.find(ntk) == wanted.end()) continue; + EvictionCandidate c; + c.key = key; + c.location = loc; + c.last_accessed_at = last_accessed; + c.size = loc.size; + result[ntk].push_back(std::move(c)); + } + } + } + + // 2. Honor the ordering hint and the per-bucket cap. For LRU, partial_sort + // is enough when a cap is set (an eviction tick is seconds, not a hot + // path). max_per_bucket == 0 means "no cap". + const auto older_first = [](const EvictionCandidate& a, const EvictionCandidate& b) { + return a.last_accessed_at < b.last_accessed_at; + }; + for (auto& [ntk, candidates] : result) { + const bool cap = max_per_bucket > 0 && candidates.size() > max_per_bucket; + if (order == EvictionOrder::kLeastRecentlyAccessed) { + if (cap) { + std::partial_sort(candidates.begin(), candidates.begin() + max_per_bucket, candidates.end(), + older_first); + candidates.resize(max_per_bucket); + } else { + std::sort(candidates.begin(), candidates.end(), older_first); + } + } else if (cap) { + candidates.resize(max_per_bucket); + } + } + return result; +} + +// ===================================================================== +// Client reads +// ===================================================================== + +std::optional InMemoryMasterMetadataStore::GetClient( + const std::string& node_id) const { + std::shared_lock lock(meta_mutex_); + auto it = clients_.find(node_id); + if (it == clients_.end()) return std::nullopt; + return it->second; +} + +bool InMemoryMasterMetadataStore::IsClientAlive(const std::string& node_id) const { + std::shared_lock lock(meta_mutex_); + return IsClientAliveLocked(node_id); +} + +std::optional InMemoryMasterMetadataStore::GetPeerAddress( + const std::string& node_id) const { + std::shared_lock lock(meta_mutex_); + auto it = clients_.find(node_id); + if (it == clients_.end()) return std::nullopt; + return it->second.peer_address; +} + +std::vector InMemoryMasterMetadataStore::ListAliveClients() const { + std::shared_lock lock(meta_mutex_); + std::vector result; + for (const auto& [id, record] : clients_) { + if (record.status == ClientStatus::ALIVE) result.push_back(record); + } + return result; +} + +std::unordered_map InMemoryMasterMetadataStore::GetAlivePeerView() const { + std::shared_lock lock(meta_mutex_); + std::unordered_map view; + view.reserve(clients_.size()); + for (const auto& [id, record] : clients_) { + if (record.status != ClientStatus::ALIVE) continue; + view.emplace(record.node_id, record.peer_address); + } + return view; +} + +std::size_t InMemoryMasterMetadataStore::AliveClientCount() const { + std::shared_lock lock(meta_mutex_); + std::size_t count = 0; + for (const auto& [id, record] : clients_) { + if (record.status == ClientStatus::ALIVE) ++count; + } + return count; +} + +std::vector InMemoryMasterMetadataStore::GetClientTags( + const std::string& node_id) const { + std::shared_lock lock(meta_mutex_); + auto it = clients_.find(node_id); + if (it == clients_.end()) return {}; + return it->second.tags; +} + +// ===================================================================== +// External-KV reads +// ===================================================================== + +std::vector InMemoryMasterMetadataStore::MatchExternalKv( + const std::vector& hashes, bool count_as_hit, + std::chrono::system_clock::time_point now) { + // count_as_hit mutates external_kv_hits_, so take meta_mutex_'s exclusive lock + // in that case; a pure read stays shared. This is the one formerly-shared path + // that becomes exclusive on meta_mutex_ (§2a), but it's one acquisition per + // RPC, not per hash. external_kv lives entirely in the meta domain, so this + // never touches a block shard. + std::unordered_map>> acc; + + auto match_into = [&]() { + for (const auto& hash : hashes) { + auto it = external_kv_entries_.find(hash); + if (it == external_kv_entries_.end()) continue; + for (const auto& [node_id, tiers] : it->second) { + auto& by_tier = acc[node_id]; + for (TierType tier : tiers) by_tier[tier].push_back(hash); + } + } + }; + + if (count_as_hit) { + std::unique_lock lock(meta_mutex_); + match_into(); + // Increment each unique matched hash once and stamp last_seen = now. + std::unordered_set matched; + for (const auto& [node_id, by_tier] : acc) { + for (const auto& [tier, hs] : by_tier) { + for (const auto& h : hs) matched.insert(h); + } + } + for (const auto& h : matched) { + auto& entry = external_kv_hits_[h]; + ++entry.count; + if (entry.last_seen < now) entry.last_seen = now; + } + } else { + std::shared_lock lock(meta_mutex_); + match_into(); + } + + std::vector result; + result.reserve(acc.size()); + for (auto& [node_id, by_tier] : acc) { + NodeMatch m; + m.node_id = node_id; + m.hashes_by_tier = std::move(by_tier); + result.push_back(std::move(m)); + } + return result; +} + +std::vector InMemoryMasterMetadataStore::GetExternalKvHitCounts( + const std::vector& hashes) const { + std::shared_lock lock(meta_mutex_); + std::vector out; + std::unordered_set seen; + seen.reserve(hashes.size()); + for (const auto& hash : hashes) { + if (!seen.insert(hash).second) continue; + auto it = external_kv_hits_.find(hash); + if (it == external_kv_hits_.end()) continue; + out.push_back(ExternalKvHitCountEntry{hash, it->second.count}); + } + return out; +} + +std::size_t InMemoryMasterMetadataStore::GetExternalKvCount(const std::string& node_id) const { + std::shared_lock lock(meta_mutex_); + std::size_t count = 0; + for (const auto& [hash, nodes] : external_kv_entries_) { + (void)hash; + if (nodes.count(node_id)) ++count; + } + return count; +} + +} // namespace mori::umbp diff --git a/src/umbp/distributed/master/master_server.cpp b/src/umbp/distributed/master/master_server.cpp index 3ebd9db5a..2d99ccc05 100644 --- a/src/umbp/distributed/master/master_server.cpp +++ b/src/umbp/distributed/master/master_server.cpp @@ -34,8 +34,9 @@ #include "mori/utils/mori_log.hpp" #include "umbp.grpc.pb.h" #include "umbp/common/env_time.h" -#include "umbp/distributed/master/external_kv_block_index.h" -#include "umbp/distributed/master/external_kv_hit_index.h" +#include "umbp/distributed/master/evict_strategy.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" +#include "umbp/distributed/master/master_metadata_store.h" #include "umbp/distributed/master/master_metrics.h" #include "umbp/distributed/routing/router.h" #include "umbp_peer.grpc.pb.h" @@ -72,16 +73,6 @@ uint32_t HitQueryMaxBatch() { return v; } -uint64_t NowNs() { - return static_cast(std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()) - .count()); -} - -uint64_t ToNs(std::chrono::seconds value) { - return static_cast(std::chrono::duration_cast(value).count()); -} - KvEvent FromProtoEvent(const ::umbp::KvEvent& pe) { KvEvent ev; switch (pe.kind()) { @@ -196,22 +187,21 @@ MasterServerConfig MasterServerConfig::FromEnvironment() { return cfg; } +// Out-of-line special members (declared in config.h). Defined here because +// this TU has the complete strategy types in scope. +MasterServerConfig::MasterServerConfig() = default; +MasterServerConfig::~MasterServerConfig() = default; +MasterServerConfig::MasterServerConfig(MasterServerConfig&&) noexcept = default; +MasterServerConfig& MasterServerConfig::operator=(MasterServerConfig&&) noexcept = default; + // --------------------------------------------------------------------------- // gRPC service implementation // --------------------------------------------------------------------------- class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Service { public: - UMBPMasterServiceImpl(ClientRegistry& registry, GlobalBlockIndex& index, - ExternalKvBlockIndex& external_kv_index, - ExternalKvHitIndex& external_kv_hit_index, Router& router, + UMBPMasterServiceImpl(IMasterMetadataStore& store, Router& router, const ClientRegistryConfig& config, mori::metrics::MetricsServer* metrics) - : registry_(registry), - index_(index), - external_kv_index_(external_kv_index), - external_kv_hit_index_(external_kv_hit_index), - router_(router), - config_(config), - metrics_(metrics) {} + : store_(store), router_(router), config_(config), metrics_(metrics) {} // -------- Client lifecycle -------- @@ -227,13 +217,21 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } const auto& engine_desc_str = request->engine_desc(); - std::vector engine_desc_bytes(engine_desc_str.begin(), engine_desc_str.end()); - - std::vector tags(request->tags().begin(), request->tags().end()); + ClientRegistration registration; + registration.node_id = request->node_id(); + registration.node_address = request->node_address(); + registration.tier_capacities = caps; + registration.peer_address = request->peer_address(); + registration.engine_desc_bytes.assign(engine_desc_str.begin(), engine_desc_str.end()); + registration.tags.assign(request->tags().begin(), request->tags().end()); + + // stale_after mirrors ClientRegistry::ExpiryDuration() (heartbeat_ttl × + // max_missed_heartbeats) so a TTL-stale ALIVE record the reaper hasn't yet + // flipped can still be re-registered (hazard #2). + const auto stale_after = config_.heartbeat_ttl * config_.max_missed_heartbeats; const bool registered = - registry_.RegisterClient(request->node_id(), request->node_address(), caps, - request->peer_address(), engine_desc_bytes, tags); + store_.RegisterClient(registration, std::chrono::system_clock::now(), stale_after); if (!registered) { return grpc::Status(grpc::StatusCode::ALREADY_EXISTS, "node is already alive and cannot be re-registered"); @@ -252,7 +250,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser grpc::Status UnregisterClient(grpc::ServerContext* /*ctx*/, const ::umbp::UnregisterClientRequest* request, ::umbp::UnregisterClientResponse* /*response*/) override { - registry_.UnregisterClient(request->node_id()); + store_.UnregisterClient(request->node_id()); UpdateClientCountMetric(); return grpc::Status::OK; } @@ -279,13 +277,65 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser bundles.push_back(std::move(bundle)); } + // §3e heartbeat adapter: the wire protocol ships EventBundle[] (each with + // its own seq) but IMasterMetadataStore::ApplyHeartbeat takes one seq at a + // time. Translate here — the store sees one seq per call, which is exactly + // what the seq-CAS (hazard #1) requires. + const auto now = std::chrono::system_clock::now(); uint64_t acked_seq = 0; bool request_full_sync = false; - auto status = - registry_.Heartbeat(request->node_id(), caps, bundles, request->is_full_sync(), - request->delta_seq_baseline(), &acked_seq, &request_full_sync); + ClientStatus client_status = ClientStatus::ALIVE; + + if (request->is_full_sync()) { + // Full sync replaces this node's locations wholesale and re-baselines + // last_applied_seq to delta_seq_baseline. Flatten every bundle's events + // into one ApplyHeartbeat call (ReplaceNodeLocations keeps only ADDs). + std::vector events; + for (const auto& bundle : bundles) { + for (const auto& ev : bundle.events) events.push_back(ev); + } + auto result = store_.ApplyHeartbeat(request->node_id(), request->delta_seq_baseline(), now, + caps, events, /*is_full_sync=*/true); + if (result.status == HeartbeatResult::UNKNOWN) { + client_status = ClientStatus::UNKNOWN; + } else { + acked_seq = result.acked_seq; + } + } else if (bundles.empty()) { + // Keepalive heartbeat: no events, but liveness must still refresh so the + // reaper doesn't expire an idle-but-alive node. seq=0 deterministically + // hits the SEQ_GAP branch, which bumps last_heartbeat + status←ALIVE + // without advancing last_applied_seq; the gap is ignored (no full-sync + // request) because there is nothing to recover. + auto result = store_.ApplyHeartbeat(request->node_id(), /*seq=*/0, now, caps, + /*events=*/{}, /*is_full_sync=*/false); + if (result.status == HeartbeatResult::UNKNOWN) { + client_status = ClientStatus::UNKNOWN; + } else { + acked_seq = result.acked_seq; + } + } else { + // Delta path: apply bundles in ascending-seq order. A real forward gap + // short-circuits the loop and requests a full sync, leaving earlier + // bundles applied (Risk item 2 — application is per-bundle, not atomic + // across the batch). + for (const auto& bundle : bundles) { + auto result = store_.ApplyHeartbeat(request->node_id(), bundle.seq, now, caps, + bundle.events, /*is_full_sync=*/false); + if (result.status == HeartbeatResult::UNKNOWN) { + client_status = ClientStatus::UNKNOWN; + break; + } + if (result.status == HeartbeatResult::SEQ_GAP) { + acked_seq = result.acked_seq; + request_full_sync = true; + break; + } + acked_seq = result.acked_seq; + } + } - response->set_status(static_cast<::umbp::ClientStatus>(status)); + response->set_status(static_cast<::umbp::ClientStatus>(client_status)); response->set_acked_seq(acked_seq); response->set_request_full_sync(request_full_sync); @@ -293,7 +343,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser if (metrics_ != nullptr && request->tier_kv_counts_size() > 0) { mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; - for (const auto& tag : registry_.GetClientTags(request->node_id())) { + for (const auto& tag : store_.GetClientTags(request->node_id())) { const auto sep = tag.find('='); if (sep != std::string::npos) { base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); @@ -466,7 +516,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser grpc::Status BatchLookup(grpc::ServerContext* /*ctx*/, const ::umbp::BatchLookupRequest* request, ::umbp::BatchLookupResponse* response) override { std::vector keys(request->keys().begin(), request->keys().end()); - auto found = index_.BatchLookupExists(keys); + auto found = store_.BatchExistsBlock(keys); for (bool b : found) response->add_found(b); return grpc::Status::OK; } @@ -484,7 +534,11 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } const TierType tier = static_cast(request->tier()); - if (!registry_.IsClientAlive(request->node_id())) { + std::vector hashes(request->hashes().begin(), request->hashes().end()); + // Alive-check + write are fused into one atomic store call (closes the + // TOCTOU gap between the old IsClientAlive check and Register — §2b.5). + const bool applied = store_.RegisterExternalKvIfAlive(request->node_id(), hashes, tier); + if (!applied) { MORI_UMBP_WARN("[Server] ReportExternalKvBlocks rejected: node not alive: {}", request->node_id()); if (metrics_) { @@ -497,14 +551,12 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser return grpc::Status::OK; } - std::vector hashes(request->hashes().begin(), request->hashes().end()); - const size_t mutated = external_kv_index_.Register(request->node_id(), hashes, tier); if (metrics_) { const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, {"tier", TierTypeName(tier)}}; metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL, MORI_UMBP_METRIC_EXT_KV_REPORT_BLOCKS_TOTAL_HELP, labels, - static_cast(mutated)); + static_cast(hashes.size())); metrics_->addCounter( MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL, MORI_UMBP_METRIC_EXT_KV_REPORT_TOTAL_HELP, {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); @@ -524,13 +576,13 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser const TierType tier = static_cast(request->tier()); std::vector hashes(request->hashes().begin(), request->hashes().end()); - const size_t mutated = external_kv_index_.Unregister(request->node_id(), hashes, tier); + store_.UnregisterExternalKv(request->node_id(), hashes, tier); if (metrics_) { const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, {"tier", TierTypeName(tier)}}; metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL_HELP, labels, - static_cast(mutated)); + static_cast(hashes.size())); metrics_->addCounter( MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); @@ -546,13 +598,12 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } const TierType tier = static_cast(request->tier()); - const size_t mutated = external_kv_index_.UnregisterByNodeAtTier(request->node_id(), tier); + // Whole-tier wipe. UnregisterExternalKvByTier returns void (the store + // interface does not surface a mutated-block count), so the per-block + // counter is no longer emitted for this admin path; the revoke_total + // counter still records the operation. + store_.UnregisterExternalKvByTier(request->node_id(), tier); if (metrics_) { - const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, - {"tier", TierTypeName(tier)}}; - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL_HELP, labels, - static_cast(mutated)); metrics_->addCounter( MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, {{"node", request->node_id()}, {"tier", TierTypeName(tier)}, {"result", "ok"}}); @@ -567,13 +618,10 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "node_id must not be empty"); } - const size_t mutated = external_kv_index_.UnregisterByNode(request->node_id()); + // All-tier wipe for one node (full-sync recovery path). Returns void, so + // the per-block counter is dropped here as in the per-tier wipe above. + store_.UnregisterExternalKvByNode(request->node_id()); if (metrics_) { - const mori::metrics::MetricsServer::Labels labels = {{"node", request->node_id()}, - {"tier", "ALL"}}; - metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL, - MORI_UMBP_METRIC_EXT_KV_REVOKE_BLOCKS_TOTAL_HELP, labels, - static_cast(mutated)); metrics_->addCounter(MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL, MORI_UMBP_METRIC_EXT_KV_REVOKE_TOTAL_HELP, {{"node", request->node_id()}, {"tier", "ALL"}, {"result", "ok"}}); @@ -585,9 +633,14 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser const ::umbp::MatchExternalKvRequest* request, ::umbp::MatchExternalKvResponse* response) override { std::vector hashes(request->hashes().begin(), request->hashes().end()); - auto matches = external_kv_index_.Match(hashes); - - auto peer_map = registry_.GetAlivePeerView(); + // The store fuses the match with the hit-count increment + last_seen stamp + // (when count_as_hit) into one lock acquisition; the handler no longer + // makes a separate IncrementHits call. `now` crosses the store boundary + // (system_clock) and feeds GarbageCollectHits. + auto matches = + store_.MatchExternalKv(hashes, request->count_as_hit(), std::chrono::system_clock::now()); + + auto peer_map = store_.GetAlivePeerView(); for (auto& m : matches) { auto* proto_match = response->add_matches(); proto_match->set_node_id(m.node_id); @@ -600,21 +653,6 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } } - if (request->count_as_hit() && !matches.empty()) { - std::unordered_set matched_hashes; - for (const auto& m : matches) { - for (const auto& [tier, hashes_in_tier] : m.hashes_by_tier) { - for (const auto& hash : hashes_in_tier) matched_hashes.insert(hash); - } - } - if (!matched_hashes.empty()) { - std::vector unique_matched; - unique_matched.reserve(matched_hashes.size()); - for (const auto& hash : matched_hashes) unique_matched.push_back(hash); - external_kv_hit_index_.IncrementHits(unique_matched, NowNs()); - } - } - size_t total_matched = 0; for (const auto& m : matches) total_matched += m.MatchedHashCount(); if (metrics_) { @@ -641,13 +679,11 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } std::vector hashes(request->hashes().begin(), request->hashes().end()); - std::vector> entries; - entries.reserve(hashes.size()); - external_kv_hit_index_.Lookup(hashes, &entries); - for (const auto& [hash, total] : entries) { + auto entries = store_.GetExternalKvHitCounts(hashes); + for (const auto& e : entries) { auto* entry = response->add_entries(); - entry->set_hash(hash); - entry->set_hit_count_total(total); + entry->set_hash(e.hash); + entry->set_hit_count_total(e.hit_count_total); } return grpc::Status::OK; } @@ -658,7 +694,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser if (!metrics_) return grpc::Status::OK; mori::metrics::MetricsServer::Labels base = {{"node", request->node_id()}}; - for (const auto& tag : registry_.GetClientTags(request->node_id())) { + for (const auto& tag : store_.GetClientTags(request->node_id())) { const auto sep = tag.find('='); if (sep != std::string::npos) { base.push_back({tag.substr(0, sep), tag.substr(sep + 1)}); @@ -696,7 +732,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser void UpdateClientCountMetric() { if (!metrics_) return; metrics_->setGauge(MORI_UMBP_METRIC_CLIENT_COUNT, MORI_UMBP_METRIC_CLIENT_COUNT_HELP, - static_cast(registry_.AliveClientCount())); + static_cast(store_.AliveClientCount())); } void UpdateClientCapacityMetrics(const std::string& node_id, @@ -724,10 +760,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser } } - ClientRegistry& registry_; - GlobalBlockIndex& index_; - ExternalKvBlockIndex& external_kv_index_; - ExternalKvHitIndex& external_kv_hit_index_; + IMasterMetadataStore& store_; Router& router_; ClientRegistryConfig config_; mori::metrics::MetricsServer* metrics_ = nullptr; @@ -738,18 +771,14 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser // --------------------------------------------------------------------------- MasterServer::MasterServer(MasterServerConfig config) : config_(std::move(config)), - index_(), - external_kv_index_(), - external_kv_hit_index_(), - registry_(config_.registry_config, index_, &external_kv_index_), - router_(index_, registry_, std::move(config_.get_strategy), std::move(config_.put_strategy)), - service_(std::make_unique(registry_, index_, external_kv_index_, - external_kv_hit_index_, router_, - config_.registry_config, nullptr)), + store_(std::make_unique()), + router_(*store_, std::move(config_.get_strategy), std::move(config_.put_strategy)), + service_(std::make_unique(*store_, router_, config_.registry_config, + nullptr)), peer_stub_pool_(std::make_unique()), - eviction_manager_(std::make_unique( - index_, registry_, config_.eviction_config, peer_stub_pool_.get(), - std::move(config_.evict_strategy))) { + eviction_manager_(std::make_unique(*store_, config_.eviction_config, + peer_stub_pool_.get(), + std::move(config_.evict_strategy))) { router_.SetLeaseDuration(config_.eviction_config.lease_duration); } @@ -767,7 +796,7 @@ void MasterServer::Run() { MORI_UMBP_INFO("[Master] Metrics server listening on port {}", config_.metrics_port); } - registry_.StartReaper(); + StartReaper(); eviction_manager_->Start(); StartHitIndexGc(); @@ -813,7 +842,7 @@ void MasterServer::Shutdown() { MORI_UMBP_INFO("[Master] Shutting down"); server_->Shutdown(deadline); } - registry_.StopReaper(); + StopReaper(); StopHitIndexGc(); } @@ -834,7 +863,7 @@ void MasterServer::StopHitIndexGc() { } void MasterServer::HitIndexGcLoop() { - const uint64_t ttl_ns = ToNs(HitIndexTtl()); + const auto ttl = HitIndexTtl(); while (hit_index_gc_running_) { { std::unique_lock lock(hit_index_gc_cv_mutex_); @@ -843,14 +872,55 @@ void MasterServer::HitIndexGcLoop() { } if (!hit_index_gc_running_) break; - const uint64_t now_ns = NowNs(); - const uint64_t cutoff_ns = now_ns > ttl_ns ? now_ns - ttl_ns : 0; - if (cutoff_ns == 0) continue; - const size_t dropped = external_kv_hit_index_.GarbageCollect(cutoff_ns); + // cutoff is a system_clock time_point now (hazard #7) so it's comparable + // to the last_seen the store stamps in MatchExternalKv(count_as_hit=true). + const auto cutoff = std::chrono::system_clock::now() - ttl; + const size_t dropped = store_->GarbageCollectHits(cutoff); if (dropped > 0) { MORI_UMBP_DEBUG("[Master] External KV hit index GC dropped {} entries", dropped); } } } +// --------------------------------------------------------------------------- +// Client-expiry reaper. Lifted from ClientRegistry — only the schedule lives +// here; the per-tick action is one store_->ExpireStaleClients(cutoff) call. +// --------------------------------------------------------------------------- +void MasterServer::StartReaper() { + bool expected = false; + if (!reaper_running_.compare_exchange_strong(expected, true)) return; + reaper_thread_ = std::thread(&MasterServer::ReaperLoop, this); + const auto expiry = + config_.registry_config.heartbeat_ttl * config_.registry_config.max_missed_heartbeats; + MORI_UMBP_INFO("[Reaper] Started (interval={}s, expiry={}s)", + config_.registry_config.reaper_interval.count(), expiry.count()); +} + +void MasterServer::StopReaper() { + bool expected = true; + if (!reaper_running_.compare_exchange_strong(expected, false)) return; + reaper_cv_.notify_one(); + if (reaper_thread_.joinable()) reaper_thread_.join(); + MORI_UMBP_INFO("[Reaper] Stopped"); +} + +void MasterServer::ReaperLoop() { + const auto expiry = + config_.registry_config.heartbeat_ttl * config_.registry_config.max_missed_heartbeats; + while (reaper_running_) { + { + std::unique_lock lock(reaper_cv_mutex_); + reaper_cv_.wait_for(lock, config_.registry_config.reaper_interval, + [this] { return !reaper_running_.load(); }); + } + if (!reaper_running_) break; + + const auto cutoff = std::chrono::system_clock::now() - expiry; + auto expired = store_->ExpireStaleClients(cutoff); + for (const auto& node_id : expired) { + MORI_UMBP_WARN("[Reaper] Expired client: {}", node_id); + } + } +} + } // namespace mori::umbp diff --git a/src/umbp/distributed/routing/route_get_strategy.cpp b/src/umbp/distributed/routing/route_get_strategy.cpp index bf59433d7..1690230f5 100644 --- a/src/umbp/distributed/routing/route_get_strategy.cpp +++ b/src/umbp/distributed/routing/route_get_strategy.cpp @@ -31,7 +31,7 @@ namespace mori::umbp { namespace { -std::string SummarizeLocations(const std::vector& locations) { +[[maybe_unused]] std::string SummarizeLocations(const std::vector& locations) { if (locations.empty()) return ""; std::ostringstream oss; bool first = true; @@ -67,19 +67,15 @@ size_t PickRandomIndex(const std::vector& indices) { return indices[dist(rng)]; } -} // namespace - -Location RandomRouteGetStrategy::Select(const std::vector& locations, - const std::string& /*node_id*/) { - if (locations.empty()) { - MORI_UMBP_WARN("[RouteGetStrategy] received empty location set; returning default Location"); - return {}; - } - +// Core selection cores shared by the single-key Select and the batched +// BatchSelect. Precondition: `locations` is non-empty (callers guarantee it). +Location PickRandomReplica(const std::vector& locations) { if (locations.size() == 1) { const auto& single = locations[0]; - MORI_UMBP_DEBUG("[RouteGetStrategy] single candidate selected node={} tier={} size={}", - single.node_id, TierTypeName(single.tier), single.size); + // NOTE: temporarily disabled — the macro args are evaluated even when the + // log level suppresses output; SummarizeLocations() dominated the hot path. + // MORI_UMBP_DEBUG("[RouteGetStrategy] single candidate selected node={} tier={} size={}", + // single.node_id, TierTypeName(single.tier), single.size); return single; } @@ -87,21 +83,15 @@ Location RandomRouteGetStrategy::Select(const std::vector& locations, std::uniform_int_distribution dist(0, locations.size() - 1); size_t choice = dist(rng); const auto& selected = locations[choice]; - MORI_UMBP_DEBUG( - "[RouteGetStrategy] {} candidates -> choice={} node={} tier={} size={}, candidates=[{}]", - locations.size(), choice, selected.node_id, TierTypeName(selected.tier), selected.size, - SummarizeLocations(locations)); + // NOTE: temporarily disabled — see above; SummarizeLocations() ran on every call. + // MORI_UMBP_DEBUG( + // "[RouteGetStrategy] {} candidates -> choice={} node={} tier={} size={}, candidates=[{}]", + // locations.size(), choice, selected.node_id, TierTypeName(selected.tier), selected.size, + // SummarizeLocations(locations)); return selected; } -Location TierPriorityRouteGetStrategy::Select(const std::vector& locations, - const std::string& /*node_id*/) { - if (locations.empty()) { - MORI_UMBP_WARN( - "[TierPriorityRouteGetStrategy] received empty location set; returning default Location"); - return {}; - } - +Location PickTierPriorityReplica(const std::vector& locations) { // Find the best (lowest-rank) tier present, then collect every replica on it // and choose one at random so load still spreads within a tier. int best_rank = TierReadRank(locations[0].tier); @@ -115,12 +105,70 @@ Location TierPriorityRouteGetStrategy::Select(const std::vector& locat size_t choice = PickRandomIndex(best_tier_indices); const auto& selected = locations[choice]; - MORI_UMBP_DEBUG( - "[TierPriorityRouteGetStrategy] {} candidates -> best_tier={} ({} replicas) choice node={} " - "tier={} size={}, candidates=[{}]", - locations.size(), TierTypeName(selected.tier), best_tier_indices.size(), selected.node_id, - TierTypeName(selected.tier), selected.size, SummarizeLocations(locations)); + // NOTE: temporarily disabled — the macro args are evaluated even when the log + // level suppresses output; SummarizeLocations() dominated the hot path. + // MORI_UMBP_DEBUG( + // "[TierPriorityRouteGetStrategy] {} candidates -> best_tier={} ({} replicas) choice node={} + // " "tier={} size={}, candidates=[{}]", locations.size(), TierTypeName(selected.tier), + // best_tier_indices.size(), selected.node_id, TierTypeName(selected.tier), selected.size, + // SummarizeLocations(locations)); return selected; } +} // namespace + +// Base default: one virtual dispatch per key. Kept so custom strategies that +// only override Select() get a working BatchSelect for free. An empty +// candidate list is left as a default Location (the caller treats it as "not +// routed") and Select() is not invoked — mirroring the router's pre-batch +// skip and avoiding a spurious empty-set WARN. +std::vector RouteGetStrategy::BatchSelect( + const std::vector>& per_key_locations, const std::string& node_id) { + std::vector out(per_key_locations.size()); + for (size_t i = 0; i < per_key_locations.size(); ++i) { + if (per_key_locations[i].empty()) continue; + out[i] = Select(per_key_locations[i], node_id); + } + return out; +} + +Location RandomRouteGetStrategy::Select(const std::vector& locations, + const std::string& /*node_id*/) { + if (locations.empty()) { + MORI_UMBP_WARN("[RouteGetStrategy] received empty location set; returning default Location"); + return {}; + } + return PickRandomReplica(locations); +} + +std::vector RandomRouteGetStrategy::BatchSelect( + const std::vector>& per_key_locations, const std::string& /*node_id*/) { + std::vector out(per_key_locations.size()); + for (size_t i = 0; i < per_key_locations.size(); ++i) { + if (per_key_locations[i].empty()) continue; // not routed; leave default + out[i] = PickRandomReplica(per_key_locations[i]); + } + return out; +} + +Location TierPriorityRouteGetStrategy::Select(const std::vector& locations, + const std::string& /*node_id*/) { + if (locations.empty()) { + MORI_UMBP_WARN( + "[TierPriorityRouteGetStrategy] received empty location set; returning default Location"); + return {}; + } + return PickTierPriorityReplica(locations); +} + +std::vector TierPriorityRouteGetStrategy::BatchSelect( + const std::vector>& per_key_locations, const std::string& /*node_id*/) { + std::vector out(per_key_locations.size()); + for (size_t i = 0; i < per_key_locations.size(); ++i) { + if (per_key_locations[i].empty()) continue; // not routed; leave default + out[i] = PickTierPriorityReplica(per_key_locations[i]); + } + return out; +} + } // namespace mori::umbp diff --git a/src/umbp/distributed/routing/router.cpp b/src/umbp/distributed/routing/router.cpp index 1f971837b..7f69e85e7 100644 --- a/src/umbp/distributed/routing/router.cpp +++ b/src/umbp/distributed/routing/router.cpp @@ -27,10 +27,9 @@ namespace mori::umbp { -Router::Router(GlobalBlockIndex& index, ClientRegistry& registry, - std::unique_ptr get_strategy, +Router::Router(IMasterMetadataStore& store, std::unique_ptr get_strategy, std::unique_ptr put_strategy) - : index_(index), registry_(registry) { + : store_(store) { // Default to tier-priority (HBM > DRAM > SSD): with the SSD cold tier live, a // random pick could route a key that also has a DRAM/HBM copy to the slow // SSD. Callers can still inject RandomRouteGetStrategy (or any other) via @@ -56,7 +55,7 @@ std::optional Router::RouteGet( std::optional Router::RoutePut( const std::string& key, const std::string& node_id, uint64_t block_size, const std::unordered_set& exclude_nodes) { - // Delegate to the batch path (size=1) so master-side dedup (BatchLookupExists) + // Delegate to the batch path (size=1) so master-side dedup (BatchExistsBlock) // and projected-capacity logic have a single home; a returned kAlreadyExists // now flows back through RoutePutResponse.outcome. auto results = BatchRoutePut({key}, node_id, {block_size}, exclude_nodes); @@ -71,8 +70,8 @@ std::vector> Router::BatchRoutePut( // the peer allocator stays the final ENOSPC arbiter. A keys/block_sizes length // mismatch is logged as a MORI ERROR and yields an all-nullopt result // (best-effort: no throw, every key reads as unroutable). - auto exists_mask = index_.BatchLookupExists(keys); - auto candidates = registry_.GetAliveClients(); + auto exists_mask = store_.BatchExistsBlock(keys); + auto candidates = store_.ListAliveClients(); return put_strategy_->SelectBatch(node_id, block_sizes, exists_mask, std::move(candidates), exclude_nodes); } @@ -83,10 +82,19 @@ std::vector> Router::BatchRouteGet( std::vector> results(keys.size()); // Lightweight membership view (no capacity): an O(nodes) copy of just the - // node->peer pairs, far cheaper than GetAliveClients' full per-node records. - auto node_to_peer = registry_.GetAlivePeerView(); + // node->peer pairs, far cheaper than ListAliveClients' full per-node records. + // Snapshot once for the whole batch — the master assumes it is stable for + // the duration of one BatchRouteGet. + auto node_to_peer = store_.GetAlivePeerView(); - auto all_locs = index_.BatchLookupForRouteGet(keys, exclude_nodes, lease_duration_); + // Unlike the old GlobalBlockIndex::BatchLookupForRouteGet (which read the + // clock internally), BatchLookupBlockForRouteGet takes an explicit `now` + // the router supplies — the timestamp now crosses the store boundary. + auto all_locs = store_.BatchLookupBlockForRouteGet( + keys, exclude_nodes, std::chrono::system_clock::now(), lease_duration_); + // One virtual dispatch for the whole batch; empty (unrouted) entries are + // skipped inside BatchSelect and left as a default Location. + auto selected_all = get_strategy_->BatchSelect(all_locs, node_id); for (size_t i = 0; i < keys.size(); ++i) { auto& locations = all_locs[i]; if (locations.empty()) { @@ -95,7 +103,7 @@ std::vector> Router::BatchRouteGet( keys[i]); continue; } - Location selected = get_strategy_->Select(locations, node_id); + Location selected = selected_all[i]; RouteGetResolution out; out.location = selected; diff --git a/src/umbp/doc/runtime-env-vars.md b/src/umbp/doc/runtime-env-vars.md index 785fb9367..c0577b84b 100644 --- a/src/umbp/doc/runtime-env-vars.md +++ b/src/umbp/doc/runtime-env-vars.md @@ -67,6 +67,7 @@ Read by the **master process** (`bin/master_main.cpp` via | `UMBP_HIT_QUERY_MAX_BATCH` | `4096` | count | Maximum hashes accepted by one `GetExternalKvHitCounts` request. Oversized requests return gRPC `INVALID_ARGUMENT`; the server does not truncate. | | `UMBP_ROUTE_PUT_SELECT_ALGO` | `most_available` | enum | Base RoutePut placement algorithm over eligible nodes (HBM before DRAM; SSD never a direct-put target). `most_available` = pick the node with the most projected free space; `random` = capacity-weighted random (probability proportional to projected `available_bytes`, never picks a node that cannot fit). Unknown value → default + one WARN. | | `UMBP_ROUTE_PUT_NODE_AFFINITY` | `none` | enum | Node-affinity bias layered on top of the base algorithm. `none` = pure base algorithm; `same` = try to place the whole batch on one node that fits the non-dedup total, else per-key sticky to the first picked node; `local` = per-key prefer the requester's local node. All three fall back to the base algorithm so affinity never makes a key fail that the base algorithm could route. Unknown value → default + one WARN. | +| `UMBP_MASTER_INDEX_SHARDS` | `32` | count | Number of independently-locked, key-hashed shards backing the block-location index inside `InMemoryMasterMetadataStore` (the block lock domain, separate from the single `meta_mutex_` that guards client records + external-KV). A heartbeat's event batch only takes the exclusive lock on the shards its keys hash into, so unrelated `RoutePut` / `BatchLookup` readers on other shards don't block behind a large apply; full-sync likewise becomes N small critical sections instead of one giant one. Read once at store construction via `std::strtol`; unset / unparseable / `< 1` → default `32` (a WARN is logged on unparseable input), clamped to a max of `4096`. `1` reproduces the old single-lock block index. Production guidance for heavy heartbeat fan-out (hundreds of clients): `64`. | ## Peer / pool client @@ -84,7 +85,6 @@ that has loaded `libmori_pybinds.so`). | `UMBP_RELEASE_LEASE_TIMEOUT_MS` | `1000` | ms | Per-attempt gRPC deadline for the best-effort `ReleaseSsdLease` RPC so a slow peer can't stall the reader. `min_allowed=1`. | | `UMBP_SSD_PREPARE_TIMEOUT_MS` | `0` | ms | Per-call gRPC deadline for `PrepareSsdRead` so a hung/slow peer can't stall the serial batch. `0` = fall back to `ssd_lease_timeout_s` (cluster-homogeneous). A timed-out / failed prepare is a hard not-served outcome (NOT retried, and never a miss). `min_allowed=0`. | | `UMBP_AUTO_FLUSH_EVENT_THRESHOLD` | `128` | count | Peer-side unshipped `KvEvent` outbox size at which a completed batch of puts auto-triggers a heartbeat flush (`FlushHeartbeat`), so the ADDs become visible at the master without waiting for the heartbeat interval or an explicit `Flush()`. Counted on `PeerDramAllocator` only (SSD events still wait for the interval). Parsed via `std::strtoull` (no WARN on bad input); unset / unparseable / `0` -> default `128`; set to a very large value to make auto-flush effectively never fire. Cached on first use in `MasterClient::SetPeerDramAllocator`. | -| `UMBP_MASTER_INDEX_SHARDS` | `32` | count | Number of independently-locked, key-hashed shards backing the master `GlobalBlockIndex`. A heartbeat's event batch only takes the exclusive lock on the shards its keys hash into, so unrelated `RoutePut` / `BatchLookup` readers on other shards don't block behind a large apply; full-sync (`ReplaceNodeLocations`) likewise becomes N small critical sections instead of one giant one. Read once at master start via `std::strtol`; unset / unparseable / `< 1` -> default `32` (a WARN is logged on unparseable input). `1` reproduces the old single-global-lock behavior. | ## SPDK proxy diff --git a/src/umbp/include/umbp/distributed/config.h b/src/umbp/include/umbp/distributed/config.h index edf2761d9..4f9f3cd84 100644 --- a/src/umbp/include/umbp/distributed/config.h +++ b/src/umbp/include/umbp/distributed/config.h @@ -124,6 +124,17 @@ struct MasterServerConfig { // ~MasterServerConfig to be instantiated in every TU that includes this // header, where those types are incomplete. static MasterServerConfig FromEnvironment(); + + // Special members are user-declared and defined out-of-line in + // master_server.cpp. This struct owns unique_ptrs to forward-declared + // strategy types (RouteGetStrategy / RoutePutStrategy / MasterEvictStrategy), + // so the destructor and move operations must be emitted in a TU where those + // types are complete — not implicitly instantiated in every includer of this + // header (which would require each to include all three strategy headers). + MasterServerConfig(); + ~MasterServerConfig(); + MasterServerConfig(MasterServerConfig&&) noexcept; + MasterServerConfig& operator=(MasterServerConfig&&) noexcept; }; struct ExportableDram { diff --git a/src/umbp/include/umbp/distributed/master/client_registry.h b/src/umbp/include/umbp/distributed/master/client_registry.h deleted file mode 100644 index 3845aae91..000000000 --- a/src/umbp/include/umbp/distributed/master/client_registry.h +++ /dev/null @@ -1,146 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/config.h" -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -class GlobalBlockIndex; -class ExternalKvBlockIndex; - -// node_id -> peer_address for every ALIVE node, for the read paths that only -// need to resolve an owning node to its peer-service address (BatchRouteGet, -// MatchExternalKv). It deliberately carries NO per-tier capacity, so building it -// is an O(nodes) copy of just these pairs, skipping the heavy per-node capacity -// map and engine desc that GetAliveClients copies. -using AlivePeerView = std::unordered_map; - -// Master-side membership ledger + heartbeat ingestion. In the -// master-as-advisor design this class no longer owns any allocator -// state; every per-tier capacity number it stores is the value the peer -// reported in its most recent heartbeat. Heartbeat is also the channel -// through which peer-shipped KvEvents reach GlobalBlockIndex. -class ClientRegistry { - public: - explicit ClientRegistry(const ClientRegistryConfig& config); - ClientRegistry(const ClientRegistryConfig& config, GlobalBlockIndex& index, - ExternalKvBlockIndex* external_kv_index = nullptr); - ~ClientRegistry(); - - ClientRegistry(const ClientRegistry&) = delete; - ClientRegistry& operator=(const ClientRegistry&) = delete; - - void SetBlockIndex(GlobalBlockIndex* index); - void SetExternalKvBlockIndex(ExternalKvBlockIndex* index); - - // --- Client lifecycle --- - - // Returns false when a live node with the same id already exists. - // Returns true for new registrations or re-registration of expired - // nodes. In the new design the only state master holds for a node is - // membership + last-reported tier capacities; the peer owns its own - // allocators. - bool RegisterClient(const std::string& node_id, const std::string& node_address, - const std::map& tier_capacities, - const std::string& peer_address = "", - const std::vector& engine_desc_bytes = {}, - const std::vector& tags = {}); - - // Drops the node from the registry and clears every index entry that - // belonged to it. - void UnregisterClient(const std::string& node_id); - - // Apply one heartbeat request. Returns the resulting status - // (UNKNOWN if the node isn't registered). On the success path: - // - tier_capacities replace the stored values unconditionally, - // - delta bundles are applied in seq order, with retransmissions skipped, - // - full-sync replaces this node's UMBP-owned locations. - ClientStatus Heartbeat(const std::string& node_id, - const std::map& tier_capacities, - const std::vector& bundles, bool is_full_sync, - uint64_t delta_seq_baseline, uint64_t* out_acked_seq, - bool* out_request_full_sync); - - // --- Queries --- - bool IsClientAlive(const std::string& node_id) const; - // Total registered nodes regardless of status. - size_t ClientCount() const; - // Count of nodes currently in ALIVE status (cheaper than GetAlivePeerView when - // only the count is needed — no map is built). Used by the client-count - // metric. - size_t AliveClientCount() const; - // Deep copy of every alive node's full record (membership + capacities + - // engine desc + tags), needed by callers that read capacity (RoutePut, - // eviction). Prefer GetAlivePeerView below when you only need node->peer. - std::vector GetAliveClients() const; - - // node_id -> peer_address for every ALIVE node, built on demand. Used by - // BatchRouteGet and MatchExternalKv; far cheaper than GetAliveClients since it - // skips the per-node capacity map / engine desc. - AlivePeerView GetAlivePeerView() const; - - // Returns the tags registered for node_id, or empty if not found. - std::vector GetClientTags(const std::string& node_id) const; - - // --- Reaper control --- - // The reaper only expires nodes whose last_heartbeat has aged past - // `heartbeat_ttl × max_missed_heartbeats`. No allocation reaper — - // pending state lives at the peer in this design. - void StartReaper(); - void StopReaper(); - - private: - ClientRegistryConfig config_; - GlobalBlockIndex* index_ = nullptr; - ExternalKvBlockIndex* external_kv_index_ = nullptr; - - mutable std::shared_mutex mutex_; - std::unordered_map clients_; - - std::thread reaper_thread_; - std::atomic reaper_running_{false}; - std::mutex reaper_cv_mutex_; - std::condition_variable reaper_cv_; - - void ReaperLoop(); - void ReapExpiredClients(); - - std::chrono::seconds ExpiryDuration() const { - return config_.heartbeat_ttl * config_.max_missed_heartbeats; - } -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/evict_strategy.h b/src/umbp/include/umbp/distributed/master/evict_strategy.h index 7e3eacf95..4da9aaec9 100644 --- a/src/umbp/include/umbp/distributed/master/evict_strategy.h +++ b/src/umbp/include/umbp/distributed/master/evict_strategy.h @@ -27,8 +27,7 @@ #include #include -#include "umbp/distributed/master/global_block_index.h" // EvictionCandidate -#include "umbp/distributed/types.h" // TierType +#include "umbp/distributed/types.h" // EvictionCandidate, TierType namespace mori::umbp { diff --git a/src/umbp/include/umbp/distributed/master/eviction_manager.h b/src/umbp/include/umbp/distributed/master/eviction_manager.h index 61c51dbe6..caa2737af 100644 --- a/src/umbp/include/umbp/distributed/master/eviction_manager.h +++ b/src/umbp/include/umbp/distributed/master/eviction_manager.h @@ -33,8 +33,7 @@ namespace mori::umbp { -class GlobalBlockIndex; -class ClientRegistry; +class IMasterMetadataStore; class MasterEvictStrategy; // Fire-and-forget callback for shipping EvictKey RPCs to a peer. The @@ -62,7 +61,7 @@ class EvictionManager { // // `strategy` is the victim-selection policy; null installs the default // LruMasterEvictStrategy (the single default-fallback site). - EvictionManager(GlobalBlockIndex& index, ClientRegistry& registry, const EvictionConfig& config, + EvictionManager(IMasterMetadataStore& store, const EvictionConfig& config, EvictKeyDispatcher* dispatcher = nullptr, std::unique_ptr strategy = nullptr); ~EvictionManager(); @@ -81,8 +80,7 @@ class EvictionManager { // background timer loop. void RunOnce(); - GlobalBlockIndex& index_; - ClientRegistry& registry_; + IMasterMetadataStore& store_; EvictionConfig config_; EvictKeyDispatcher* dispatcher_; std::unique_ptr strategy_; diff --git a/src/umbp/include/umbp/distributed/master/external_kv_block_index.h b/src/umbp/include/umbp/distributed/master/external_kv_block_index.h deleted file mode 100644 index fa0161e0f..000000000 --- a/src/umbp/include/umbp/distributed/master/external_kv_block_index.h +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -/// Lightweight index for externally-managed KV blocks (e.g. sglang HiCache). -/// Each (node, hash) pair tracks the set of tiers the node has reported. -class ExternalKvBlockIndex { - public: - ExternalKvBlockIndex() = default; - ~ExternalKvBlockIndex() = default; - - ExternalKvBlockIndex(const ExternalKvBlockIndex&) = delete; - ExternalKvBlockIndex& operator=(const ExternalKvBlockIndex&) = delete; - - // Mutators return the count of actually changed (hash, node, tier) tuples. - size_t Register(const std::string& node_id, const std::vector& hashes, - TierType tier); - size_t Unregister(const std::string& node_id, const std::vector& hashes, - TierType tier); - size_t UnregisterByNodeAtTier(const std::string& node_id, TierType tier); - size_t UnregisterByNode(const std::string& node_id); - - struct NodeMatch { - std::string node_id; - std::map> hashes_by_tier; - - size_t MatchedHashCount() const { - std::unordered_set seen; - for (const auto& [tier, hashes] : hashes_by_tier) { - for (const auto& h : hashes) seen.insert(h); - } - return seen.size(); - } - }; - - std::vector Match(const std::vector& hashes) const; - size_t GetKvCount(const std::string& node_id) const; - - private: - mutable std::shared_mutex mutex_; - std::unordered_map>> entries_; -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/external_kv_hit_index.h b/src/umbp/include/umbp/distributed/master/external_kv_hit_index.h deleted file mode 100644 index 4f4bc7720..000000000 --- a/src/umbp/include/umbp/distributed/master/external_kv_hit_index.h +++ /dev/null @@ -1,81 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace mori::umbp { - -// Per-hash cumulative hit counter for external KV placement matches. -// Entries are created only for hashes that were actually matched by -// MatchExternalKv(count_as_hit=true); revoke does not remove them. -class ExternalKvHitIndex { - public: - static constexpr size_t kShards = 256; - - ExternalKvHitIndex() = default; - ~ExternalKvHitIndex() = default; - - ExternalKvHitIndex(const ExternalKvHitIndex&) = delete; - ExternalKvHitIndex& operator=(const ExternalKvHitIndex&) = delete; - - // Caller owns request-level de-duplication. Each input hash is incremented - // exactly once by this call. - void IncrementHits(const std::vector& unique_hashes, uint64_t now_ns); - - // Sparse lookup. Missing hashes are skipped, and duplicate query hashes - // produce at most one output entry. - size_t Lookup(const std::vector& hashes, - std::vector>* out) const; - - // Drop entries whose last activity is older than cutoff_ns. - size_t GarbageCollect(uint64_t cutoff_ns); - - size_t Size() const; - - private: - struct Entry { - std::atomic total{0}; - std::atomic last_seen_ns{0}; - }; - - struct Shard { - mutable std::shared_mutex mu; - std::unordered_map> entries; - }; - - static size_t ShardIdx(std::string_view hash); - static void UpdateLastSeen(Entry* entry, uint64_t now_ns); - - std::array shards_; -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/global_block_index.h b/src/umbp/include/umbp/distributed/master/global_block_index.h deleted file mode 100644 index bb127454d..000000000 --- a/src/umbp/include/umbp/distributed/master/global_block_index.h +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -struct BlockEntry { - std::vector locations; - BlockMetrics metrics; - - std::atomic lease_expiry_rep{0}; - std::atomic last_accessed_rep{0}; - std::atomic atomic_access_count{0}; - - void GrantLease(std::chrono::steady_clock::duration duration) { - GrantLease(std::chrono::steady_clock::now(), duration); - } - - // Overload taking a caller-supplied "now": lets a batch capture the clock - // once and reuse it across many keys instead of re-reading it per key. - void GrantLease(std::chrono::steady_clock::time_point now, - std::chrono::steady_clock::duration duration) { - auto expiry = now + duration; - lease_expiry_rep.store(expiry.time_since_epoch().count(), std::memory_order_release); - } - - bool IsLeased() const { - auto now_rep = std::chrono::steady_clock::now().time_since_epoch().count(); - return lease_expiry_rep.load(std::memory_order_acquire) > now_rep; - } - - void RecordAccessAtomic() { RecordAccessAtomic(std::chrono::steady_clock::now()); } - - // Overload taking a caller-supplied "now"; see GrantLease above. - void RecordAccessAtomic(std::chrono::steady_clock::time_point now) { - last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); - atomic_access_count.fetch_add(1, std::memory_order_relaxed); - } - - std::chrono::steady_clock::time_point GetLastAccessed() const { - auto rep = last_accessed_rep.load(std::memory_order_acquire); - return std::chrono::steady_clock::time_point(std::chrono::steady_clock::duration(rep)); - } -}; - -struct EvictionCandidate { - std::string key; - Location location; - std::chrono::steady_clock::time_point last_accessed_at; - uint64_t size; -}; - -// Master-side projection of every peer's owned-key set. In the -// master-as-advisor design this index is *only* mutated through the -// event-shipping heartbeat — there are no per-Put or per-Eviction -// master RPCs. Routing and eviction read from here. -class GlobalBlockIndex { - public: - // The forward + reverse index is split into `UMBP_MASTER_INDEX_SHARDS` - // independently-locked shards (key-hashed). A heartbeat's event batch then - // only takes the exclusive lock on the shards its keys land in, so unrelated - // route/lookup readers on other shards never block behind a large apply. - GlobalBlockIndex(); - ~GlobalBlockIndex() = default; - - GlobalBlockIndex(const GlobalBlockIndex&) = delete; - GlobalBlockIndex& operator=(const GlobalBlockIndex&) = delete; - - // --- Mutators (event-driven only) --- - - // Apply one peer's heartbeat-shipped event batch. Returns the count - // of location mutations. ADD with a (node_id, tier) that already exists - // for the key is an idempotent no-op on the location's size. - // REMOVE for an unknown (key, node_id, tier) is a silent no-op. - size_t ApplyEvents(const std::string& node_id, const std::vector& events); - - // Replace this node's full set of locations with the ADDs carried in `adds`. - void ReplaceNodeLocations(const std::string& node_id, const std::vector& adds); - - void RemoveByNode(const std::string& node_id); - - // Bump last_accessed_at and access_count. Lock-free under the shared lock. - void RecordAccess(const std::string& key); - - // Grant a time-limited lease to protect a key from eviction. - void GrantLease(const std::string& key, std::chrono::steady_clock::duration duration); - - // Batched Lookup + filter + (on non-empty result) RecordAccess + GrantLease, - // under a single shared_lock. - std::vector> BatchLookupForRouteGet( - const std::vector& keys, const std::unordered_set& exclude_nodes, - std::chrono::steady_clock::duration lease_duration); - - // --- Queries --- - - std::vector Lookup(const std::string& key) const; - - // Batched existence check — single shared_lock acquisition for the - // whole batch. Read-only, no access-count or lease side-effects. - // Returns a vector parallel to `keys` where entry i is true iff the - // key has at least one registered Location. - std::vector BatchLookupExists(const std::vector& keys) const; - - std::optional GetMetrics(const std::string& key) const; - - // --- Eviction --- - - struct NodeTierKey { - std::string node_id; - TierType tier; - bool operator<(const NodeTierKey& o) const { - if (node_id != o.node_id) return node_id < o.node_id; - return tier < o.tier; - } - bool operator==(const NodeTierKey& o) const { return node_id == o.node_id && tier == o.tier; } - }; - - std::vector FindEvictionCandidates( - const std::set& overloaded_node_tiers) const; - - private: - // One independently-locked partition of the key space. A key lives in - // exactly one shard (chosen by hash), so its forward entry and its - // reverse-index membership are always co-located under the same lock. - struct Shard { - mutable std::shared_mutex mutex; - std::unordered_map entries; - // Reverse index (shard-local): node_id -> set of this shard's keys the - // node owns. Lets ReplaceNodeLocations/RemoveByNode skip a full scan. - std::unordered_map> node_to_keys; - }; - - size_t shard_index(std::string_view key) const { - return std::hash{}(key) % num_shards_; - } - Shard& shard_at(size_t i) const { return *shards_[i]; } - Shard& shard_for(std::string_view key) const { return shard_at(shard_index(key)); } - - const size_t num_shards_; - std::vector> shards_; -}; - -} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h b/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h new file mode 100644 index 000000000..20f1118be --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/in_memory_master_metadata_store.h @@ -0,0 +1,240 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// In-process implementation of IMasterMetadataStore. +// +// This is the single-master / unit-test backend: it folds the four former +// state holders (GlobalBlockIndex, ClientRegistry, ExternalKvBlockIndex, +// ExternalKvHitIndex) into one class. The logic is a near-verbatim lift of +// those classes; what changes is the locking and that every timestamp crossing +// the interface boundary is now caller-supplied system_clock — see the hazards +// in master_metadata_store.h. +// +// Two lock domains (NOT one global mutex): +// - meta_mutex_ (shared_mutex) guards client records, external-kv locations, +// and per-hash hit counts (former ClientRegistry / ExternalKvBlockIndex / +// ExternalKvHitIndex state). +// - The block index (former GlobalBlockIndex) is split into +// UMBP_MASTER_INDEX_SHARDS key-hashed shards, each with its own mutex; a +// key's forward entry and its reverse-index membership share a shard. A +// heartbeat's events only lock the shards they hash into, so readers on +// other shards never block behind a large apply — this restores PR #440's +// per-shard heartbeat concurrency that a single mutex would serialize away. +// Cross-domain writes (ApplyHeartbeat / UnregisterClient / ExpireStaleClients) +// run the meta section then the block section, meta first, WITHOUT nesting the +// locks — atomic within each domain but NOT globally atomic. See the atomicity +// contract in master_metadata_store.h. +// +// Per-hash hit counts live in process memory and are lost on restart, exactly +// as the old ExternalKvHitIndex did; crash-durability is a Redis-backend +// concern only. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/master_metadata_store.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +class InMemoryMasterMetadataStore : public IMasterMetadataStore { + public: + // Reads UMBP_MASTER_INDEX_SHARDS once and builds that many block shards + // (default 32, clamped to [1, 4096]; 1 reproduces the old single-lock block + // index). See IndexShardCount() in the .cpp. + InMemoryMasterMetadataStore(); + ~InMemoryMasterMetadataStore() override = default; + + InMemoryMasterMetadataStore(const InMemoryMasterMetadataStore&) = delete; + InMemoryMasterMetadataStore& operator=(const InMemoryMasterMetadataStore&) = delete; + + // --- Cross-store writes --- + bool RegisterClient(const ClientRegistration& registration, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after) override; + void UnregisterClient(const std::string& node_id) override; + HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, + std::chrono::system_clock::time_point now, + const std::map& caps, + const std::vector& events, bool is_full_sync) override; + std::vector ExpireStaleClients( + std::chrono::system_clock::time_point cutoff) override; + + // --- External-KV writes --- + bool RegisterExternalKvIfAlive(const std::string& node_id, const std::vector& hashes, + TierType tier) override; + void UnregisterExternalKv(const std::string& node_id, const std::vector& hashes, + TierType tier) override; + void UnregisterExternalKvByTier(const std::string& node_id, TierType tier) override; + void UnregisterExternalKvByNode(const std::string& node_id) override; + std::size_t GarbageCollectHits(std::chrono::system_clock::time_point cutoff) override; + + // --- Block reads --- + std::vector LookupBlock(const std::string& key) const override; + std::vector LookupBlockForRouteGet( + const std::string& key, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) override; + std::vector> BatchLookupBlockForRouteGet( + const std::vector& keys, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) override; + std::vector BatchExistsBlock(const std::vector& keys) const override; + std::map> EnumerateEvictionCandidates( + const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now) const override; + + // --- Client reads --- + std::optional GetClient(const std::string& node_id) const override; + bool IsClientAlive(const std::string& node_id) const override; + std::optional GetPeerAddress(const std::string& node_id) const override; + std::vector ListAliveClients() const override; + std::unordered_map GetAlivePeerView() const override; + std::size_t AliveClientCount() const override; + std::vector GetClientTags(const std::string& node_id) const override; + + // --- External-KV reads --- + std::vector MatchExternalKv(const std::vector& hashes, bool count_as_hit, + std::chrono::system_clock::time_point now) override; + std::vector GetExternalKvHitCounts( + const std::vector& hashes) const override; + std::size_t GetExternalKvCount(const std::string& node_id) const override; + + private: + // One block's locations + LRU/lease metadata. Lifted from + // GlobalBlockIndex::BlockEntry, but the lease/access mutators now take a + // caller-supplied `now` (system_clock) instead of reading the clock + // internally — the value crosses the store boundary (hazard #7). The + // lease/access state stays in atomics so the RouteGet path can mutate it + // under the shared lock of the key's shard, exactly as today (§2a). + struct BlockEntry { + std::vector locations; + BlockMetrics metrics; + + std::atomic lease_expiry_rep{0}; + std::atomic last_accessed_rep{0}; + std::atomic atomic_access_count{0}; + + void GrantLease(std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration duration) { + auto expiry = now + duration; + lease_expiry_rep.store(expiry.time_since_epoch().count(), std::memory_order_release); + } + + bool IsLeased(std::chrono::system_clock::time_point now) const { + return lease_expiry_rep.load(std::memory_order_acquire) > now.time_since_epoch().count(); + } + + void RecordAccessAtomic(std::chrono::system_clock::time_point now) { + last_accessed_rep.store(now.time_since_epoch().count(), std::memory_order_release); + atomic_access_count.fetch_add(1, std::memory_order_relaxed); + } + + std::chrono::system_clock::time_point GetLastAccessed() const { + auto rep = last_accessed_rep.load(std::memory_order_acquire); + return std::chrono::system_clock::time_point(std::chrono::system_clock::duration(rep)); + } + }; + + // Per-hash cumulative hit counter (lifted from ExternalKvHitIndex, collapsed + // from 256 atomic shards to a single map under meta_mutex_). last_seen is + // system_clock now that it crosses the boundary and feeds GarbageCollectHits. + struct HitEntry { + uint64_t count = 0; + std::chrono::system_clock::time_point last_seen; + }; + + // One independently-locked partition of the block key space. A key lives in + // exactly one shard (chosen by hash), so its forward entry and its + // reverse-index membership are always co-located under the same lock. + struct Shard { + mutable std::shared_mutex mutex; + // Block locations (from GlobalBlockIndex). + std::unordered_map entries; + // Reverse index (shard-local) node_id -> this shard's keys the node owns, + // so node-scoped removal skips a full scan. + std::unordered_map> node_to_keys; + }; + + size_t shard_index(std::string_view key) const { + return std::hash{}(key) % num_shards_; + } + Shard& shard_at(size_t i) const { return *shards_[i]; } + Shard& shard_for(std::string_view key) const { return shard_at(shard_index(key)); } + + // --- Block-index writes: acquire the relevant shard lock(s) internally. + // Callers MUST NOT hold meta_mutex_ when calling these (see the cross-domain + // write contract in master_metadata_store.h) --- + size_t ApplyEventsToShards(const std::string& node_id, const std::vector& events, + std::chrono::system_clock::time_point now); + void ReplaceNodeLocationsInShards(const std::string& node_id, const std::vector& adds, + std::chrono::system_clock::time_point now); + void RemoveBlocksByNodeInShards(const std::string& node_id); + + // --- Single-shard block helpers (caller MUST hold that shard's unique lock). + // Static members so they can name the private Shard / BlockEntry types. --- + // Apply one ADD/REMOVE to a shard; returns 1 iff a location was mutated (a + // duplicate ADD is an idempotent no-op that does NOT count). + static size_t ApplyAddOrRemoveLocked(Shard& shard, const std::string& node_id, const KvEvent& ev, + std::chrono::system_clock::time_point now); + // Drop every location owned by node_id in one shard, driven by the shard-local + // reverse index (O(node's keys in this shard), not a full scan). + static void RemoveNodeLocationsLocked(Shard& shard, const std::string& node_id); + + // --- Meta-domain locked helpers (caller MUST hold meta_mutex_'s unique lock + // for the mutators, shared for IsClientAliveLocked) --- + void RemoveExternalKvByNodeLocked(const std::string& node_id); + bool IsClientAliveLocked(const std::string& node_id) const; + + // Meta domain: client records + external-kv locations + per-hash hit counts. + mutable std::shared_mutex meta_mutex_; + + // Client records (from ClientRegistry). + std::unordered_map clients_; + + // External-KV locations (from ExternalKvBlockIndex): hash -> node -> tier-set. + // Keyed hash-first so MatchExternalKv (the hot RPC path) stays O(1) per hash. + std::unordered_map>> + external_kv_entries_; + + // Per-hash hit counts (from ExternalKvHitIndex). + std::unordered_map external_kv_hits_; + + // Block domain: key-hashed shards, each independently locked. + const size_t num_shards_; + std::vector> shards_; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/master_metadata_store.h b/src/umbp/include/umbp/distributed/master/master_metadata_store.h new file mode 100644 index 000000000..7f963fb90 --- /dev/null +++ b/src/umbp/include/umbp/distributed/master/master_metadata_store.h @@ -0,0 +1,510 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Master-side metadata storage interface for mori::umbp. +// +// Goal: make the master stateless so multiple master replicas can serve +// traffic concurrently (HA). Today GlobalBlockIndex, ClientRegistry, +// ExternalKvBlockIndex, and ExternalKvHitIndex hold their state in +// in-process unordered_maps guarded by std::shared_mutex — that's a +// split-brain hazard the moment you run more than one master process, +// because each replica's view of locations, liveness, leases, LRU, and +// hit counts drifts independently. The fix is to pull ALL durable AND +// volatile master state out behind one abstract interface — +// IMasterMetadataStore — and ship it to a shared backend. +// +// Concrete target backend is Redis (Cluster); a SQL backend is NOT a +// target. The read path does one store round-trip per RouteGet +// (lookup + lease + access in one Lua script) which is fine against +// Redis but ruinous against an OLTP database. If a SQL backend is +// ever needed it'd be for cold archival, not the hot path served by +// this interface. +// +// ===================================================================== +// Why one interface and not three (BlockLocations / ClientRecords / +// ExternalKv) per the earlier sketch: +// ===================================================================== +// - Every cross-store write today touches 2 or 3 of those stores: +// * UnregisterClient → registry + block index + external_kv +// * ReapExpiredClients → registry + block index + external_kv +// * Heartbeat → registry seq-CAS + block index events +// * RegisterExternalKvBlocks → registry alive-check + external_kv +// * RegisterClient re-registration → registry read TTL + write +// - Splitting forces a portable cross-store transaction abstraction. +// In-memory implements it trivially (shared_mutex). Redis does +// not: cross-keyspace MULTI/EXEC or Lua only works if all three +// namespaces hash to the same slot, which leaks "must run on the +// same cluster shard" up into the API. +// - Collapsing to one interface lets each implementation choose its +// own atomicity primitive instead of inventing a portable +// transaction abstraction. +// - Hot-path reads stay easy to identify because they're grouped in +// the read section below; internally an in-memory impl can still +// keep separate sub-maps for code organization. The interface just +// stops pretending they're independently swappable. +// +// ===================================================================== +// What does NOT move behind this interface: +// ===================================================================== +// - EvictionManager policy (watermark math, victim grouping, EvictKey +// RPC dispatch). The store returns LRU-ordered candidates limited +// by a byte budget; the manager decides how to spend that budget +// and ships the RPCs. Stateless under HA — each tick's decision is +// a pure function of the store's snapshot, so concurrent eviction +// passes on different replicas converge instead of fighting. +// - Reaper loop scheduling (timer + cv). Only the per-pass DB action +// moves down, via ExpireStaleClients. The schedule is per-replica +// with no shared state; the ExpireStaleClients call is idempotent +// so multiple replicas can safely run reaper passes concurrently. +// - Hit-index GC loop scheduling (timer + cv). Only the per-pass +// action moves down, via GarbageCollectHits. Same per-replica / +// idempotent reasoning as the reaper. +// +// In particular, lease_expiry / last_accessed_at / access_count and the +// per-hash hit counts DO move into the store — see hazards #4 and #7. +// +// ===================================================================== +// Critical design decisions / hazards (carried forward from review): +// ===================================================================== +// 1. Heartbeat is a CAS, not Get-then-Update. +// The current in-memory ClientRegistry::Heartbeat reads +// last_applied_seq, decides whether to accept, and writes the new +// seq + caps + last_heartbeat under one unique_lock. If this is +// split into Get() then UpdateHeartbeatState() across an external +// backend, two concurrent heartbeats for the same node can both +// observe seq N, both decide "in order," and one silently +// corrupts last_applied_seq. ApplyHeartbeat MUST take seq as the +// CAS value and atomically check `seq == last_applied_seq + 1` +// inside the implementation (one Lua script on Redis; +// shared_mutex unique_lock for the in-memory impl). See +// HeartbeatResult below. +// +// 2. RegisterClient must accept TTL-stale ALIVE rows. +// The current RegisterClient (client_registry.cpp:85-91) allows +// re-registration when `(now - last_heartbeat > ExpiryDuration())` +// even if status==ALIVE (i.e. the reaper hasn't flipped it yet). +// A naive InsertIfNotAlive would reject. The new RegisterClient +// therefore takes a `stale_after` duration so the implementation +// enforces this in the same atomic step. +// +// 3. EXPIRED rows are KEPT, not erased. +// Today's ReapExpiredClients erases rows from the map. +// ExpireStaleClients flips status ALIVE → EXPIRED and keeps the +// row so a re-registration can replace an EXPIRED record cleanly. +// Behavioral change vs. today; consumers must filter status when +// counting (use AliveClientCount, not "size of GetClient over all +// ids", which doesn't exist anyway). +// +// 4. lease_expiry, last_accessed_at, access_count are IN the store. +// Earlier sketches kept them in a process-local LeaseAccessTracker +// to avoid per-RouteGet store writes; that's only safe with a +// single master process. Under HA-stateless masters, two replicas +// can hold conflicting in-memory leases for the same key — +// replica A's lease doesn't block replica B's eviction loop — +// and LRU views diverge per replica. Both are correctness bugs, +// not stale-state inconveniences. So the tracker is gone: +// RouteGet hits the store via LookupBlockForRouteGet which +// atomically reads locations, sets lease_expiry, and bumps +// last_accessed_at / access_count in one Lua script. The +// BatchLookupBlockForRouteGet variant amortizes the cost over +// the prefix-match path the router already uses. +// +// 5. BatchLookupBlockForRouteGet exists because router.cpp:139 +// issues N Lookup() calls in a hot loop. N round-trips against +// Redis is unacceptable; one batched Lua call is required. +// +// 6. Sync vs async. Methods are synchronous; concurrency comes from +// the gRPC handler threads. Cleaner than future-returning every +// method. +// +// 7. Persisted timestamps are system_clock, not steady_clock. +// steady_clock has a per-process arbitrary epoch — its values +// are not meaningful to a different process or after master +// restart, so they cannot live in a Redis row. Every timestamp +// that crosses the IMasterMetadataStore boundary — registration +// times, heartbeat timestamps, lease_expiry, last_accessed_at, +// reaper cutoff, and the hit-count last_seen — is therefore +// system_clock::time_point. No steady_clock anywhere below this +// interface. Assumes NTP-disciplined clocks across master +// replicas: a backward wall-clock jump effectively grants longer +// leases until the clock recovers. +// +// 8. `last_acked_seq` is gone from the heartbeat path. +// The current ClientRegistry::Heartbeat takes a `last_acked_seq` +// parameter (client_registry.cpp:138) but its body ignores it — +// the master gap-checks against its own last_applied_seq. The new +// interface drops the parameter rather than passing through a +// value no implementation will read. The peer wire still carries +// last_acked_seq for the peer's own ack-on-progress logic; the +// master_server adapter simply doesn't forward it down here. +// +// NodeTierKey, NodeMatch, ClientRegistration, HeartbeatResult, and +// EvictionCandidate are defined in umbp/distributed/types.h — they are +// part of this store's contract. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +// ===================================================================== +// IMasterMetadataStore — single durable-state interface for the master. +// ===================================================================== +// +// All methods are thread-safe. The master's gRPC handler thread pool +// calls into a single shared instance from many threads concurrently — +// implementations must provide their own synchronization (shared_mutex +// for in-memory, single-script Lua atomicity for Redis). Callers do +// not add external locking around these calls. +// +// Every write method below is atomic in isolation. The methods that span +// what used to be multiple stores (UnregisterClient, ApplyHeartbeat, +// ExpireStaleClients, RegisterExternalKvIfAlive, and the hit-counting branch +// of MatchExternalKv) are why the four stores were merged behind one +// interface. HOW atomic they are across those former boundaries is +// backend-specific: +// +// - RedisMasterMetadataStore: globally atomic — one Lua script over the +// shared hash tag mutates all keyspaces in one step (also what makes a +// large full_sync atomic; see TODO(payload-sizing) below). +// +// - InMemoryMasterMetadataStore: atomic within each lock domain, NOT +// globally atomic across them. Two domains: meta_mutex_ (client records + +// external-kv + hit counts) and N key-hashed block shards. Cross-domain +// writes run the meta section first, then the block section AFTER +// releasing meta_mutex_ — never nested. This trades global atomicity for +// concurrency: heartbeats to different shards apply in parallel instead of +// serializing on one lock (PR #440's headline win). Reader-visible windows: +// * ApplyHeartbeat: seq/caps advanced before the events are visible. +// * Unregister/Expire: client gone/EXPIRED while its block locations +// briefly linger; conversely, since the block wipe runs after +// meta_mutex_ is released, a same-node re-register + full_sync inside +// the window can have its fresh locations dropped by the late wipe. +// * full_sync: atomic per shard (a key's clear+replay never tears), not +// across shards. +// All benign: a lingering location for a non-ALIVE node is filtered out by +// GetAlivePeerView (router.cpp), so it never routes to a dead node; the +// re-register case is mere under-representation (a cache miss a re-put +// heals). Same windows main/#440 had with its two independent locks. (A +// same-node heartbeat reorder is unreachable: the client serializes its +// heartbeats via hb_send_mutex_ and the RPC returns only after apply.) +// +// TODO(atomicity-contract): DECIDED for in-memory = approach A above. Open for +// Redis: confirm the Lua path keeps this method set globally atomic. The +// conformance test (test_in_memory_master_metadata_store.cpp) races a reader +// against an in-flight full_sync/heartbeat and asserts RouteGet sees +// old-or-new locations (never torn) and never resolves a peer for an +// unregistered node. +// +// Expected implementations: +// - InMemoryMasterMetadataStore: meta_mutex_ + N key-hashed block +// shards (see the atomicity contract above). Mostly a mechanical +// lift of the former classes, kept as private helpers. Used for +// single-master deployments and unit tests. The per-hash hit counts +// live in process memory and are lost on restart, exactly as the +// current ExternalKvHitIndex does. +// - RedisMasterMetadataStore: cross-keyspace ops via Lua scripts. +// All key namespaces (node:, block:, extkv:, hit:, lru:, lease:) +// must share a hash tag (e.g. `{umbp:}:node:`) +// so they hash to the same slot on Redis Cluster — that's what +// makes cross-namespace Lua atomic, and what makes the hit counts +// crash-durable. +class IMasterMetadataStore { + public: + virtual ~IMasterMetadataStore() = default; + + // =================================================================== + // Cross-store write operations — each call is atomic. + // =================================================================== + + // CAS-style registration. Inserts a fresh ALIVE record. + // - Returns true on new registration. + // - Returns true and replaces the record if an EXPIRED record for + // the same node_id exists. + // - Returns true and replaces the record if an ALIVE record exists + // whose last_heartbeat is older than `stale_after` — handles the + // "reaper hasn't run yet but the record is TTL-stale" case that + // today's RegisterClient permits (see hazard #2 in header + // preamble). + // - Returns false if an ALIVE non-stale record already exists. + // `now` is supplied by the caller so tests can inject time. Uses + // system_clock because the value is persisted in the backend; see + // hazard #7 in header preamble. The store sets last_heartbeat and + // registered_at to `now`, status to ALIVE, and last_applied_seq to 0 + // — the caller does not (and cannot) populate those fields. + // In production, callers derive `stale_after` from + // ClientRegistryConfig::ExpiryDuration() (heartbeat_ttl × + // max_missed_heartbeats); tests inject their own value. + virtual bool RegisterClient(const ClientRegistration& registration, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after) = 0; + + // Drop the client from the client store AND drop every block location + // belonging to it AND drop every external-kv entry belonging to it. + // Idempotent on missing clients. + virtual void UnregisterClient(const std::string& node_id) = 0; + + // Heartbeat ingestion. Atomically: + // 1. Looks up the client record; returns UNKNOWN if absent (no + // other state touched). + // 2. If !is_full_sync and seq != last_applied_seq + 1, returns + // SEQ_GAP with acked_seq = last_applied_seq. SEQ_GAP still + // bumps last_heartbeat and sets status←ALIVE so the reaper + // doesn't kill a node that's heartbeating but mid-recovery; + // caps and last_applied_seq are NOT touched on SEQ_GAP. The + // seq check is the CAS that keeps the gap check race-free + // under concurrent heartbeats — DO NOT implement as a + // separate Get() then UpdateHeartbeatState(); two in-flight + // heartbeats can both observe the same seq and corrupt + // last_applied_seq (see hazard #1 in header preamble). + // 3. On APPLIED, updates caps, last_heartbeat, last_applied_seq, + // status←ALIVE. + // 4. Applies events to block locations: + // * is_full_sync=true → replace every location for node_id + // with the ADDs in events; REMOVE entries ignored. + // * is_full_sync=false → ADD with existing (node,tier) + // overwrites size; REMOVE for unknown (key,node,tier) is + // a silent no-op. + // Returns APPLIED with acked_seq = seq on success. + // + // TODO(payload-sizing): the `events` vector is bounded only by what + // the peer chooses to ship in one heartbeat batch. A full_sync from + // a peer with millions of keys produces a single Lua script of that + // size, which blocks every other Redis client (Redis is + // single-threaded). The contract should be: implementations MAY + // chunk internally for !is_full_sync, but is_full_sync MUST apply + // atomically (a half-applied ReplaceNodeLocations would leave the + // index in a torn state). That in turn forces an upper bound on + // peer-side full_sync batch size — decide and document the cap + // (e.g. 100k events) here, and have the peer fragment larger + // resyncs into multiple full_sync calls or shift to a snapshot- + // then-delta protocol before the Redis backend ships. + virtual HeartbeatResult ApplyHeartbeat(const std::string& node_id, uint64_t seq, + std::chrono::system_clock::time_point now, + const std::map& caps, + const std::vector& events, bool is_full_sync) = 0; + + // Reaper pass. Atomically: + // - Flips status ALIVE → EXPIRED for every record with + // last_heartbeat < cutoff. EXPIRED records are KEPT in the store, + // not erased — see hazard #3 in header preamble. + // - Drops every block location belonging to those clients. + // - Drops every external-kv entry belonging to those clients. + // Returns the affected node_ids for logging. + virtual std::vector ExpireStaleClients( + std::chrono::system_clock::time_point cutoff) = 0; + + // =================================================================== + // External-KV writes — alive-check + mutation atomic together. + // =================================================================== + + // Add `tier` to the tier-set of every (node_id, hash). Idempotent: + // re-registering at the same tier is a no-op; registering at a new + // tier adds a bucket without touching existing tiers. + // Returns true if the alive-check passed and the writes were applied + // (even if every write was a no-op because the entries already existed). + // Returns false if node_id was not ALIVE and nothing was written, so + // the caller can meter the reject without the impl having to log it. + virtual bool RegisterExternalKvIfAlive(const std::string& node_id, + const std::vector& hashes, TierType tier) = 0; + + // Remove `tier` from the tier-set of every (node_id, hash). Other + // tiers for the same hash untouched. (node,hash) entry dropped when + // its tier-set becomes empty. Does NOT check liveness — peers may + // ship unregister during teardown after status has flipped. + virtual void UnregisterExternalKv(const std::string& node_id, + const std::vector& hashes, TierType tier) = 0; + + // Remove `tier` from every hash registered by `node_id` (whole-tier + // wipe — admin path, not heartbeat). + virtual void UnregisterExternalKvByTier(const std::string& node_id, TierType tier) = 0; + + // Drop every external-kv entry (all tiers) belonging to `node_id` without + // touching the client record. Backs the live RevokeAllExternalKvBlocksForNode + // RPC, which a peer issues to wipe its external-KV registration before a + // full re-sync. Does NOT check liveness. Idempotent on unknown nodes. + virtual void UnregisterExternalKvByNode(const std::string& node_id) = 0; + + // Drop every per-hash hit-count entry whose last_seen < cutoff. + // Returns the number of entries dropped. Replaces + // ExternalKvHitIndex::GarbageCollect; the cutoff is a system_clock + // time_point (not a uint64_t ns) because last_seen now crosses the + // store boundary — hazard #7. Called by the master's hit-index GC + // loop on each tick with cutoff = system_clock::now() - max_age. + virtual std::size_t GarbageCollectHits(std::chrono::system_clock::time_point cutoff) = 0; + + // =================================================================== + // Reads. None require cross-store atomicity; implementations SHOULD + // make each one a single backend round-trip. + // =================================================================== + + // --- Block locations --- + + // Plain location lookup. Returns every location for `key` without + // granting a lease or recording an access. Pure read — no side + // effects on lease_expiry, last_accessed_at, or access_count. + virtual std::vector LookupBlock(const std::string& key) const = 0; + + // RouteGet primitive. Atomically reads every location for `key`, + // filters out locations whose node_id is in `exclude_nodes`, + // and — only if at least one location survives the filter — sets + // lease_expiry to now + lease_duration and bumps last_accessed_at + // to now / access_count by 1. Returns the filtered locations, or + // empty if the key has no locations or all were excluded. + // Filtering inside the store (not post-hoc in the caller) is + // required so that fully-excluded keys do not receive a lease or + // an access bump — granting those would perturb LRU ordering and + // extend eviction protection for keys the caller explicitly chose + // to skip. Splitting this into separate Lookup / GrantLease / + // RecordAccess methods would be three round trips per RouteGet + // AND would not be atomic across master replicas — see hazard #4. + virtual std::vector LookupBlockForRouteGet( + const std::string& key, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) = 0; + + // Vectorized RouteGet primitive. Same per-key semantics as + // LookupBlockForRouteGet (leases granted and access recorded only + // for keys that have at least one non-excluded location). Result + // parallel to `keys`; absent or fully-excluded keys yield empty + // inner vectors. One round trip for the whole batch — + // router.cpp:139 today issues N Lookup() calls in a hot loop and + // this is the single-RTT replacement (see hazard #5). + virtual std::vector> BatchLookupBlockForRouteGet( + const std::vector& keys, const std::unordered_set& exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration) = 0; + + // Batched existence — pure read, no lease grant, no access record. + // Used by BatchRoutePut for dedup (router.cpp:112): a writer + // landing on an existing key isn't a "read" and must not extend + // the lease or perturb LRU ordering. One round trip. + virtual std::vector BatchExistsBlock(const std::vector& keys) const = 0; + + // Policy-neutral eviction-candidate enumeration. For each (node, tier) + // in `buckets`, returns the eviction-eligible rows (lease_expiry <= now + // filtered out), in the order requested by `order`, capped at + // `max_per_bucket` rows per bucket (0 means no cap). Result map is keyed + // by the same NodeTierKeys as the input; absent buckets had no eligible + // candidates. Taking the whole bucket set in one call lets a single Lua + // script fan out over every overloaded bucket in one round trip — + // important when dozens of (node, tier) pairs are over watermark. + // + // This call does NOT make the eviction decision: it neither sees the + // byte budget nor trims to it. The caller (EvictionManager) hands the + // returned candidates to a MasterEvictStrategy, which owns victim + // selection. `order` and `max_per_bucket` are purely a performance + // affordance — they let a backend with an index serve the cheapest + // rows instead of shipping every candidate; the eviction policy lives + // in the strategy, not here. + // + // No EraseBlock on this interface — peers ship REMOVEs on their + // next heartbeat after EvictKey executes, so the only mutation + // channels for block locations are ApplyHeartbeat / + // UnregisterClient / ExpireStaleClients. + // + // How `order` is honored is an implementation detail: the in-memory + // backend does a full entries_ scan + partial_sort per tick (an + // eviction tick is seconds, not a hot path), while the Redis backend + // can maintain a per-(node, tier) ZSET keyed by last_accessed_at and + // serve a ZRANGE. The contract is only "return eligible candidates in + // the requested order, capped," not the index mechanism. + virtual std::map> EnumerateEvictionCandidates( + const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now) const = 0; + + // --- Client records --- + + // Returns the record regardless of status (ALIVE or EXPIRED). Caller + // filters when needed. + virtual std::optional GetClient(const std::string& node_id) const = 0; + + // Hot-path liveness check. Exists as its own method so a Redis backend + // can answer it with a single status field read instead of fetching the + // whole ClientRecord just to filter on status. + virtual bool IsClientAlive(const std::string& node_id) const = 0; + + // Single-node peer-address lookup. Exists as its own method so a + // Redis backend can answer with a single HGET on the node hash + // instead of fetching the whole ClientRecord just to read + // peer_address. The legacy router linear-scans GetAliveClients() + // per RouteGet for the same value; GetPeerAddress replaces that + // with one read. Returns std::nullopt for unknown node_id; + // EXPIRED records still surface their peer_address. + virtual std::optional GetPeerAddress(const std::string& node_id) const = 0; + + // ALIVE only — does not include EXPIRED records. + virtual std::vector ListAliveClients() const = 0; + + // Lightweight membership view: node_id -> peer_address for every ALIVE + // client, and nothing else. Exists as its own method so the hot RouteGet / + // MatchExternalKv paths can snapshot just the peer addresses they need + // instead of copying every full ClientRecord (capacities, engine desc, + // tags) via ListAliveClients. A Redis backend can serve this from a single + // node->peer projection (or an SMEMBERS + per-node HGET peer) rather than + // HGETALL per node; it stays a single logical round trip and needs no + // cross-keyspace atomicity. EXPIRED clients are excluded. + virtual std::unordered_map GetAlivePeerView() const = 0; + + virtual std::size_t AliveClientCount() const = 0; + + virtual std::vector GetClientTags(const std::string& node_id) const = 0; + + // --- External KV --- + + // Returns matches grouped by node WITHOUT peer_address. Callers + // that need peer addresses join with ListAliveClients() (snapshot + // once per response) or GetPeerAddress(node_id) per-node; + // embedding peer_address in NodeMatch would force every + // implementation to read from two namespaces on every match. + // When `count_as_hit` is true, atomically increments the per-hash + // hit counter for every matched hash AND stamps that hash's + // last_seen = `now`, all in one lock acquisition (lookup + + // increment + stamp). When false, pure read — no hit counts + // touched and `now` is ignored. `now` is system_clock because + // last_seen is persisted / feeds GarbageCollectHits (hazard #7). + virtual std::vector MatchExternalKv(const std::vector& hashes, + bool count_as_hit, + std::chrono::system_clock::time_point now) = 0; + + // Sparse per-hash hit-count read. Returns an entry for each requested + // hash that has a recorded count (hashes with no recorded hits may be + // omitted). Replaces ExternalKvHitIndex::Lookup; backs the live + // GetExternalKvHitCounts RPC. Pure read. + virtual std::vector GetExternalKvHitCounts( + const std::vector& hashes) const = 0; + + virtual std::size_t GetExternalKvCount(const std::string& node_id) const = 0; +}; + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/master_server.h b/src/umbp/include/umbp/distributed/master/master_server.h index d382478ac..1874a300d 100644 --- a/src/umbp/include/umbp/distributed/master/master_server.h +++ b/src/umbp/include/umbp/distributed/master/master_server.h @@ -34,12 +34,9 @@ #include "mori/metrics/prometheus_metrics_server.hpp" #include "umbp/distributed/config.h" -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/master/evict_strategy.h" #include "umbp/distributed/master/eviction_manager.h" -#include "umbp/distributed/master/external_kv_block_index.h" -#include "umbp/distributed/master/external_kv_hit_index.h" -#include "umbp/distributed/master/global_block_index.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" +#include "umbp/distributed/master/master_metadata_store.h" #include "umbp/distributed/routing/route_get_strategy.h" #include "umbp/distributed/routing/route_put_strategy.h" #include "umbp/distributed/routing/router.h" @@ -64,10 +61,10 @@ class MasterServer { private: MasterServerConfig config_; - GlobalBlockIndex index_; - ExternalKvBlockIndex external_kv_index_; - ExternalKvHitIndex external_kv_hit_index_; - ClientRegistry registry_; + // Single owner of all master metadata state (block locations, client + // records, external-KV locations, hit counts). Declared before router_ and + // service_ so it outlives the references they hold. + std::unique_ptr store_; Router router_; std::unique_ptr metrics_server_; @@ -94,6 +91,19 @@ class MasterServer { std::atomic hit_index_gc_running_{false}; std::mutex hit_index_gc_cv_mutex_; std::condition_variable hit_index_gc_cv_; + + // Client-expiry reaper. Formerly owned by ClientRegistry; the schedule + // (timer + cv) lives here now and the per-tick action is a single + // store_->ExpireStaleClients(cutoff) call, where cutoff is on the + // system_clock basis so it's comparable to the records' last_heartbeat. + void StartReaper(); + void StopReaper(); + void ReaperLoop(); + + std::thread reaper_thread_; + std::atomic reaper_running_{false}; + std::mutex reaper_cv_mutex_; + std::condition_variable reaper_cv_; }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/routing/route_get_strategy.h b/src/umbp/include/umbp/distributed/routing/route_get_strategy.h index 52201f666..2d9c55bd2 100644 --- a/src/umbp/include/umbp/distributed/routing/route_get_strategy.h +++ b/src/umbp/include/umbp/distributed/routing/route_get_strategy.h @@ -39,6 +39,21 @@ class RouteGetStrategy { /// @param node_id The requesting client's node_id (for locality-aware strategies). /// @return The chosen Location to read from. virtual Location Select(const std::vector& locations, const std::string& node_id) = 0; + + /// Batched form of Select: choose one replica per key in a single virtual + /// call. The router uses this on the BatchRouteGet hot path so a batch of N + /// keys costs one virtual dispatch instead of N — the shipped strategies + /// override it with a tight internal loop that avoids per-key dispatch. + /// @param per_key_locations One candidate list per key, in request order. + /// An empty list means the key has no routable replica: Select is NOT + /// invoked for it and the corresponding output is a default-constructed + /// Location the caller must treat as "not routed". + /// @param node_id The requesting client's node_id. + /// @return One Location per input entry (same size and order as input). + /// The default implementation loops over Select(); subclasses that only + /// implement Select() keep working unchanged. + virtual std::vector BatchSelect( + const std::vector>& per_key_locations, const std::string& node_id); }; /// Default strategy: uniform random selection among replicas. @@ -46,6 +61,8 @@ class RouteGetStrategy { class RandomRouteGetStrategy : public RouteGetStrategy { public: Location Select(const std::vector& locations, const std::string& node_id) override; + std::vector BatchSelect(const std::vector>& per_key_locations, + const std::string& node_id) override; }; /// Tier-priority strategy: prefer the fastest tier present (HBM > DRAM > SSD), @@ -55,6 +72,8 @@ class RandomRouteGetStrategy : public RouteGetStrategy { class TierPriorityRouteGetStrategy : public RouteGetStrategy { public: Location Select(const std::vector& locations, const std::string& node_id) override; + std::vector BatchSelect(const std::vector>& per_key_locations, + const std::string& node_id) override; }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/routing/router.h b/src/umbp/include/umbp/distributed/routing/router.h index 7b359d530..2d89431c9 100644 --- a/src/umbp/include/umbp/distributed/routing/router.h +++ b/src/umbp/include/umbp/distributed/routing/router.h @@ -29,8 +29,7 @@ #include #include -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/master/global_block_index.h" +#include "umbp/distributed/master/master_metadata_store.h" #include "umbp/distributed/routing/route_get_strategy.h" #include "umbp/distributed/routing/route_put_strategy.h" @@ -45,8 +44,7 @@ struct RouteGetResolution { class Router { public: - Router(GlobalBlockIndex& index, ClientRegistry& registry, - std::unique_ptr get_strategy = nullptr, + Router(IMasterMetadataStore& store, std::unique_ptr get_strategy = nullptr, std::unique_ptr put_strategy = nullptr); ~Router() = default; @@ -75,14 +73,13 @@ class Router { const std::vector& keys, const std::string& node_id, const std::unordered_set& exclude_nodes); - void SetLeaseDuration(std::chrono::steady_clock::duration d) { lease_duration_ = d; } + void SetLeaseDuration(std::chrono::system_clock::duration d) { lease_duration_ = d; } private: - GlobalBlockIndex& index_; - ClientRegistry& registry_; + IMasterMetadataStore& store_; std::unique_ptr get_strategy_; std::unique_ptr put_strategy_; - std::chrono::steady_clock::duration lease_duration_{std::chrono::seconds{10}}; + std::chrono::system_clock::duration lease_duration_{std::chrono::seconds{10}}; }; } // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/types.h b/src/umbp/include/umbp/distributed/types.h index 118dc232b..ebc6a5da3 100644 --- a/src/umbp/include/umbp/distributed/types.h +++ b/src/umbp/include/umbp/distributed/types.h @@ -26,6 +26,8 @@ #include #include #include +#include +#include #include namespace mori::umbp { @@ -61,6 +63,38 @@ struct Location { } }; +// Identifies one (node, tier) capacity bucket — the granularity at which +// eviction budgets and overload are tracked. Hoisted from +// GlobalBlockIndex because it is part of the master metadata store contract. +struct NodeTierKey { + std::string node_id; + TierType tier; + bool operator<(const NodeTierKey& o) const { + if (node_id != o.node_id) return node_id < o.node_id; + return tier < o.tier; + } + bool operator==(const NodeTierKey& o) const { return node_id == o.node_id && tier == o.tier; } +}; + +// One node's external-KV match result, grouped by tier. A single hash may +// appear in MORE THAN ONE tier bucket when a node holds multiple physical +// copies (e.g. HBM + DRAM mirror). std::map iterates in sorted TierType +// order, so the first non-empty bucket is the fastest available tier. +// Hoisted from ExternalKvBlockIndex because it is part of the master +// metadata store contract. +struct NodeMatch { + std::string node_id; + std::map> hashes_by_tier; + + size_t MatchedHashCount() const { + std::unordered_set seen; + for (const auto& [tier, hashes] : hashes_by_tier) { + for (const auto& h : hashes) seen.insert(h); + } + return seen.size(); + } +}; + enum class ClientStatus : int { UNKNOWN = 0, ALIVE = 1, @@ -68,11 +102,32 @@ enum class ClientStatus : int { }; struct BlockMetrics { - std::chrono::steady_clock::time_point created_at; - std::chrono::steady_clock::time_point last_accessed_at; + std::chrono::system_clock::time_point created_at; + std::chrono::system_clock::time_point last_accessed_at; uint64_t access_count = 0; }; +// One eviction-eligible (key, location) row returned by the master metadata +// store's candidate enumeration. Hoisted from GlobalBlockIndex because it is +// part of the IMasterMetadataStore contract (EnumerateEvictionCandidates +// returns these; MasterEvictStrategy consumes them). +struct EvictionCandidate { + std::string key; + Location location; + std::chrono::system_clock::time_point last_accessed_at; + uint64_t size; +}; + +// Ordering hint for IMasterMetadataStore::EnumerateEvictionCandidates. This is +// a performance affordance, NOT eviction policy: it only tells the store what +// order to return rows in so a backend with an index (e.g. a Redis ZSET keyed +// by last_accessed_at) can serve the cheapest top-N rows instead of shipping +// every candidate. The actual victim decision belongs to MasterEvictStrategy. +enum class EvictionOrder : int { + kNone = 0, // no ordering guarantee; store returns in any order + kLeastRecentlyAccessed = 1, // oldest last_accessed_at first +}; + // Structured form of one (buffer_index, page_index) slot. Used by the // peer DRAM/HBM allocator to describe which page slot a write should // land in, and by ResolveKey responses to tell readers where to RDMA @@ -127,8 +182,8 @@ struct ClientRecord { std::string node_id; std::string node_address; ClientStatus status = ClientStatus::UNKNOWN; - std::chrono::steady_clock::time_point last_heartbeat; - std::chrono::steady_clock::time_point registered_at; + std::chrono::system_clock::time_point last_heartbeat; + std::chrono::system_clock::time_point registered_at; std::map tier_capacities; std::string peer_address; @@ -142,6 +197,31 @@ struct ClientRecord { std::vector tags; }; +// Input to IMasterMetadataStore::RegisterClient. Deliberately omits +// last_heartbeat / registered_at / status / last_applied_seq: those are owned +// by the store and derived from the `now` argument the caller passes alongside +// this struct. Keeping them off the input removes the "did the caller bother to +// set these?" ambiguity. +struct ClientRegistration { + std::string node_id; + std::string node_address; + std::map tier_capacities; + std::string peer_address; + std::vector engine_desc_bytes; + std::vector tags; +}; + +// Result of IMasterMetadataStore::ApplyHeartbeat. APPLIED = events accepted, +// registry updated, acked_seq advanced to the request's seq. SEQ_GAP = peer's +// seq is not last_applied_seq + 1; caller responds with a full-sync request and +// acked_seq echoes the previously applied seq so the peer reships. UNKNOWN = no +// record for node_id (peer must re-register). +struct HeartbeatResult { + enum Status { APPLIED, SEQ_GAP, UNKNOWN }; + Status status; + uint64_t acked_seq; // meaningful for APPLIED and SEQ_GAP +}; + // Helpers for logging inline const char* TierTypeName(TierType t) { switch (t) { diff --git a/src/umbp/tests/CMakeLists.txt b/src/umbp/tests/CMakeLists.txt deleted file mode 100644 index 8886e7f5d..000000000 --- a/src/umbp/tests/CMakeLists.txt +++ /dev/null @@ -1,219 +0,0 @@ -# --------------------------------------------------------------------------- -# UMBP unit tests — requires GTest and umbp_common -# --------------------------------------------------------------------------- -cmake_minimum_required(VERSION 3.14) - -include(FetchContent) - -FetchContent_Declare( - googletest - GIT_REPOSITORY https://github.com/google/googletest.git - GIT_TAG v1.14.0) - -# Prevent GoogleTest from overriding compiler/linker options when built as a -# subproject. -set(gtest_force_shared_crt - ON - CACHE BOOL "" FORCE) -FetchContent_MakeAvailable(googletest) - -enable_testing() - -include(GoogleTest) - -# --------------------------------------------------------------------------- -# test_external_kv_block_index -# --------------------------------------------------------------------------- -add_executable(test_external_kv_block_index test_external_kv_block_index.cpp) - -target_link_libraries(test_external_kv_block_index PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_external_kv_block_index PRIVATE cxx_std_17) - -gtest_discover_tests(test_external_kv_block_index) - -# --------------------------------------------------------------------------- -# test_client_registry — membership ledger: register/re-register, capacity -# round-trip, heartbeat status, and the silent-node reaper. -# --------------------------------------------------------------------------- -add_executable(test_client_registry test_client_registry.cpp) - -target_link_libraries(test_client_registry PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_client_registry PRIVATE cxx_std_17) - -gtest_discover_tests(test_client_registry) - -# --------------------------------------------------------------------------- -# test_client_registry_external_kv -# --------------------------------------------------------------------------- -add_executable(test_client_registry_external_kv - test_client_registry_external_kv.cpp) - -target_link_libraries(test_client_registry_external_kv - PRIVATE umbp_common GTest::gtest_main) - -target_compile_features(test_client_registry_external_kv PRIVATE cxx_std_17) - -gtest_discover_tests(test_client_registry_external_kv) - -# --------------------------------------------------------------------------- -# test_external_kv_hit_index -# --------------------------------------------------------------------------- -add_executable(test_external_kv_hit_index test_external_kv_hit_index.cpp) - -target_link_libraries(test_external_kv_hit_index PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_external_kv_hit_index PRIVATE cxx_std_17) - -gtest_discover_tests(test_external_kv_hit_index) - -# --------------------------------------------------------------------------- -# test_peer_dram_allocator -# --------------------------------------------------------------------------- -add_executable(test_peer_dram_allocator test_peer_dram_allocator.cpp) - -target_link_libraries(test_peer_dram_allocator PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_peer_dram_allocator PRIVATE cxx_std_17) - -gtest_discover_tests(test_peer_dram_allocator) - -# --------------------------------------------------------------------------- -# test_global_block_index_events -# --------------------------------------------------------------------------- -add_executable(test_global_block_index_events - test_global_block_index_events.cpp) - -target_link_libraries(test_global_block_index_events PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_global_block_index_events PRIVATE cxx_std_17) - -gtest_discover_tests(test_global_block_index_events) - -# --------------------------------------------------------------------------- -# test_router_dedup — master-side BatchRoutePut dedup via GlobalBlockIndex -# --------------------------------------------------------------------------- -add_executable(test_router_dedup test_router_dedup.cpp) - -target_link_libraries(test_router_dedup PRIVATE umbp_common GTest::gtest_main) - -target_compile_features(test_router_dedup PRIVATE cxx_std_17) - -gtest_discover_tests(test_router_dedup) - -# --------------------------------------------------------------------------- -# test_peer_ssd_manager — SSD tier ownership + owned-location source -# --------------------------------------------------------------------------- -add_executable(test_peer_ssd_manager test_peer_ssd_manager.cpp) - -target_link_libraries(test_peer_ssd_manager PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_peer_ssd_manager PRIVATE cxx_std_17) - -gtest_discover_tests(test_peer_ssd_manager) - -# --------------------------------------------------------------------------- -# test_peer_ssd_eviction — LRU + watermark eviction + in-flight guard -# --------------------------------------------------------------------------- -add_executable(test_peer_ssd_eviction test_peer_ssd_eviction.cpp) - -target_link_libraries(test_peer_ssd_eviction PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_peer_ssd_eviction PRIVATE cxx_std_17) - -gtest_discover_tests(test_peer_ssd_eviction) - -# --------------------------------------------------------------------------- -# test_ssd_copy_pipeline — DramCopyPin + async copy-on-commit pipeline -# --------------------------------------------------------------------------- -add_executable(test_ssd_copy_pipeline test_ssd_copy_pipeline.cpp) - -target_link_libraries(test_ssd_copy_pipeline PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_ssd_copy_pipeline PRIVATE cxx_std_17) - -gtest_discover_tests(test_ssd_copy_pipeline) - -# --------------------------------------------------------------------------- -# test_tier_priority_route_get — RouteGet tier-priority strategy -# --------------------------------------------------------------------------- -add_executable(test_tier_priority_route_get test_tier_priority_route_get.cpp) - -target_link_libraries(test_tier_priority_route_get PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_tier_priority_route_get PRIVATE cxx_std_17) - -gtest_discover_tests(test_tier_priority_route_get) - -# --------------------------------------------------------------------------- -# test_peer_ssd_read_rpc — SSD read RPC over a gRPC loopback (status -# distinctions: OK / NOT_FOUND / NO_SLOT / SIZE_TOO_LARGE + slot lifecycle). -# Needs the gRPC peer service + generated stubs (umbp_core) and protobuf/gRPC. -# --------------------------------------------------------------------------- -add_executable(test_peer_ssd_read_rpc test_peer_ssd_read_rpc.cpp) -if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") - target_link_options(test_peer_ssd_read_rpc PRIVATE -Wl,--no-as-needed) -endif() - -target_link_libraries( - test_peer_ssd_read_rpc PRIVATE umbp_core umbp_common ${_PROTOBUF_LIB} - ${_GRPCPP_LIB} GTest::gtest_main) - -target_compile_features(test_peer_ssd_read_rpc PRIVATE cxx_std_17) - -gtest_discover_tests(test_peer_ssd_read_rpc) - -# --------------------------------------------------------------------------- -# test_batch_resolve_codec — struct-of-arrays BatchResolveKeysResponse wire -# codec (batch_resolve_codec.h): encode/decode round-trip, page flattening, desc -# dedup / omit_descs, malformed-response rejection, and the payload-size win. -# Needs the generated proto (carried by umbp_common) + protobuf. -# --------------------------------------------------------------------------- -add_executable(test_batch_resolve_codec test_batch_resolve_codec.cpp) -if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") - target_link_options(test_batch_resolve_codec PRIVATE -Wl,--no-as-needed) -endif() - -target_link_libraries(test_batch_resolve_codec - PRIVATE umbp_common ${_PROTOBUF_LIB} GTest::gtest_main) - -target_compile_features(test_batch_resolve_codec PRIVATE cxx_std_17) - -gtest_discover_tests(test_batch_resolve_codec) - -# --------------------------------------------------------------------------- -# test_ssd_read_lease_gating — pure reader-side lease gating decision logic -# (ssd_read_lease.h). Header-only under test (no gRPC / RDMA deps). -# --------------------------------------------------------------------------- -add_executable(test_ssd_read_lease_gating test_ssd_read_lease_gating.cpp) - -target_link_libraries(test_ssd_read_lease_gating PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_ssd_read_lease_gating PRIVATE cxx_std_17) - -gtest_discover_tests(test_ssd_read_lease_gating) - -# --------------------------------------------------------------------------- -# test_ssd_reliability — cross-component reliability: owned-source DRAM+SSD -# merge, SSD evict -> REMOVE -> master index convergence, tier-priority over the -# real index, crash-restart discard, and observability counters. -# --------------------------------------------------------------------------- -add_executable(test_ssd_reliability test_ssd_reliability.cpp) - -target_link_libraries(test_ssd_reliability PRIVATE umbp_common - GTest::gtest_main) - -target_compile_features(test_ssd_reliability PRIVATE cxx_std_17) - -gtest_discover_tests(test_ssd_reliability) diff --git a/src/umbp/tests/test_client_registry.cpp b/src/umbp/tests/test_client_registry.cpp deleted file mode 100644 index c01a322ec..000000000 --- a/src/umbp/tests/test_client_registry.cpp +++ /dev/null @@ -1,334 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -// -// Membership-ledger unit tests for ClientRegistry: registration / re- -// registration semantics, capacity round-trip, heartbeat status, and the -// background reaper that expires silent nodes. These exercise the registry -// in isolation (no GlobalBlockIndex / RPC), complementing the index- and -// external-kv-focused suites. In the master-as-advisor design the registry -// stores only membership + the capacities a peer last reported, so the -// assertions below check reported values verbatim rather than any allocator- -// derived view. -#include - -#include -#include -#include -#include -#include - -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/types.h" - -namespace mori::umbp { -namespace { - -std::map Caps(uint64_t total, uint64_t available) { - return {{TierType::HBM, TierCapacity{total, available}}}; -} - -const ClientRecord* FindClient(const std::vector& clients, const std::string& id) { - for (const auto& c : clients) { - if (c.node_id == id) return &c; - } - return nullptr; -} - -// Drive the current 7-arg Heartbeat with no events — the membership-keepalive -// path the reaper cares about. -ClientStatus Beat(ClientRegistry& registry, const std::string& node_id, - const std::map& caps) { - uint64_t acked = 0; - bool need_full = false; - return registry.Heartbeat(node_id, caps, /*bundles=*/{}, /*is_full_sync=*/false, - /*delta_seq_baseline=*/0, &acked, &need_full); -} - -template -bool WaitUntil(Predicate&& predicate, std::chrono::milliseconds timeout, - std::chrono::milliseconds poll = std::chrono::milliseconds(100)) { - const auto deadline = std::chrono::steady_clock::now() + timeout; - while (std::chrono::steady_clock::now() < deadline) { - if (predicate()) return true; - std::this_thread::sleep_for(poll); - } - return predicate(); -} - -// heartbeat_ttl * max_missed_heartbeats == 1s, so a node ages out ~1s after -// its last heartbeat. reaper_interval keeps the sweep responsive. -ClientRegistryConfig FastExpiryConfig() { - ClientRegistryConfig config; - config.heartbeat_ttl = std::chrono::seconds(1); - config.max_missed_heartbeats = 1; - config.reaper_interval = std::chrono::seconds(1); - return config; -} - -} // namespace - -// --- Registration / membership ---------------------------------------------- - -TEST(ClientRegistryTest, RegisterSingle) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("node-1", "127.0.0.1:8080", Caps(80, 64))); - EXPECT_EQ(registry.ClientCount(), 1u); - EXPECT_TRUE(registry.IsClientAlive("node-1")); -} - -TEST(ClientRegistryTest, RegisterMultiple) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "127.0.0.1:1001", Caps(100, 90))); - EXPECT_TRUE(registry.RegisterClient("c2", "127.0.0.1:1002", Caps(110, 80))); - EXPECT_TRUE(registry.RegisterClient("c3", "127.0.0.1:1003", Caps(120, 70))); - - EXPECT_EQ(registry.ClientCount(), 3u); - EXPECT_TRUE(registry.IsClientAlive("c1")); - EXPECT_TRUE(registry.IsClientAlive("c2")); - EXPECT_TRUE(registry.IsClientAlive("c3")); -} - -TEST(ClientRegistryTest, GetAliveClientsReportsMembershipAndCapacities) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "host-a:8080", Caps(80, 64))); - EXPECT_TRUE(registry.RegisterClient("c2", "host-b:8080", Caps(96, 32))); - - const auto clients = registry.GetAliveClients(); - ASSERT_EQ(clients.size(), 2u); - - const ClientRecord* c1 = FindClient(clients, "c1"); - const ClientRecord* c2 = FindClient(clients, "c2"); - ASSERT_NE(c1, nullptr); - ASSERT_NE(c2, nullptr); - - EXPECT_EQ(c1->node_address, "host-a:8080"); - EXPECT_EQ(c2->node_address, "host-b:8080"); - EXPECT_EQ(c1->status, ClientStatus::ALIVE); - EXPECT_EQ(c2->status, ClientStatus::ALIVE); - - // Master stores the peer-reported capacities verbatim. - ASSERT_TRUE(c1->tier_capacities.count(TierType::HBM) > 0); - ASSERT_TRUE(c2->tier_capacities.count(TierType::HBM) > 0); - EXPECT_EQ(c1->tier_capacities.at(TierType::HBM).total_bytes, 80u); - EXPECT_EQ(c1->tier_capacities.at(TierType::HBM).available_bytes, 64u); - EXPECT_EQ(c2->tier_capacities.at(TierType::HBM).available_bytes, 32u); -} - -// The lightweight peer view maps node->peer (no capacity) and reflects -// membership changes. -TEST(ClientRegistryTest, AlivePeerViewTracksMembership) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "host-a:8080", Caps(80, 64), - /*peer_address=*/"peer-a:9000")); - - auto v1 = registry.GetAlivePeerView(); - EXPECT_EQ(v1.size(), 1u); - ASSERT_EQ(v1.count("c1"), 1u); - EXPECT_EQ(v1.at("c1"), "peer-a:9000"); - - // Membership change -> reflected in a freshly built view. - EXPECT_TRUE(registry.RegisterClient("c2", "host-b:8080", Caps(96, 32), - /*peer_address=*/"peer-b:9000")); - auto v2 = registry.GetAlivePeerView(); - EXPECT_EQ(v2.size(), 2u); - EXPECT_EQ(v2.at("c2"), "peer-b:9000"); -} - -// The peer view carries no capacity, so a capacity-only heartbeat leaves its -// contents (node -> peer) unchanged. -TEST(ClientRegistryTest, PeerViewIgnoresCapacity) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64), /*peer_address=*/"peer-a")); - - auto p1 = registry.GetAlivePeerView(); - EXPECT_EQ(Beat(registry, "c1", Caps(80, 8)), ClientStatus::ALIVE); - auto p2 = registry.GetAlivePeerView(); - EXPECT_EQ(p1, p2); -} - -// AliveClientCount counts only ALIVE nodes and tracks membership changes. -TEST(ClientRegistryTest, AliveClientCountTracksMembership) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_EQ(registry.AliveClientCount(), 0u); - - EXPECT_TRUE(registry.RegisterClient("c1", "addr-1", Caps(80, 64))); - EXPECT_TRUE(registry.RegisterClient("c2", "addr-2", Caps(96, 32))); - EXPECT_EQ(registry.AliveClientCount(), 2u); - - registry.UnregisterClient("c1"); - EXPECT_EQ(registry.AliveClientCount(), 1u); -} - -TEST(ClientRegistryTest, ReRegisterAliveRejected) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr-1", Caps(80, 64))); - // A live node may not silently take over its own id with a new address. - EXPECT_FALSE(registry.RegisterClient("c1", "addr-2", Caps(80, 32))); - - EXPECT_EQ(registry.ClientCount(), 1u); - const auto clients = registry.GetAliveClients(); - ASSERT_EQ(clients.size(), 1u); - EXPECT_EQ(clients[0].node_address, "addr-1"); // original record untouched -} - -TEST(ClientRegistryTest, ReRegisterExpiredAllowed) { - // No reaper here: the aged-out branch in RegisterClient (now - last_heartbeat - // > expiry) must accept the re-registration on its own. - ClientRegistry registry(FastExpiryConfig()); - EXPECT_TRUE(registry.RegisterClient("c1", "addr-1", Caps(80, 64))); - - const bool reregistered = - WaitUntil([®istry] { return registry.RegisterClient("c1", "addr-2", Caps(80, 32)); }, - std::chrono::seconds(5)); - EXPECT_TRUE(reregistered); - - EXPECT_EQ(registry.ClientCount(), 1u); - const auto clients = registry.GetAliveClients(); - ASSERT_EQ(clients.size(), 1u); - EXPECT_EQ(clients[0].node_address, "addr-2"); // new address wins - EXPECT_EQ(clients[0].status, ClientStatus::ALIVE); -} - -// --- Unregister -------------------------------------------------------------- - -TEST(ClientRegistryTest, UnregisterExisting) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - - registry.UnregisterClient("c1"); - EXPECT_EQ(registry.ClientCount(), 0u); - EXPECT_FALSE(registry.IsClientAlive("c1")); -} - -TEST(ClientRegistryTest, UnregisterUnknownIsNoop) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - - registry.UnregisterClient("nonexistent"); - EXPECT_EQ(registry.ClientCount(), 1u); - EXPECT_TRUE(registry.IsClientAlive("c1")); -} - -TEST(ClientRegistryTest, UnregisterTwiceIsSafe) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - - registry.UnregisterClient("c1"); - registry.UnregisterClient("c1"); - EXPECT_EQ(registry.ClientCount(), 0u); -} - -// --- Heartbeat --------------------------------------------------------------- - -TEST(ClientRegistryTest, HeartbeatAliveReplacesCapacities) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - - EXPECT_EQ(Beat(registry, "c1", Caps(80, 16)), ClientStatus::ALIVE); - EXPECT_TRUE(registry.IsClientAlive("c1")); - - const auto clients = registry.GetAliveClients(); - ASSERT_EQ(clients.size(), 1u); - ASSERT_TRUE(clients[0].tier_capacities.count(TierType::HBM) > 0); - // The most recent heartbeat's capacities replace the stored values. - EXPECT_EQ(clients[0].tier_capacities.at(TierType::HBM).available_bytes, 16u); -} - -TEST(ClientRegistryTest, HeartbeatUnknownReturnsUnknown) { - ClientRegistry registry(ClientRegistryConfig{}); - EXPECT_EQ(Beat(registry, "nonexistent", Caps(80, 48)), ClientStatus::UNKNOWN); -} - -// --- Reaper ------------------------------------------------------------------ - -TEST(ClientRegistryTest, ReaperExpiresIdleClient) { - ClientRegistry registry(FastExpiryConfig()); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - registry.StartReaper(); - - const bool reaped = - WaitUntil([®istry] { return registry.ClientCount() == 0; }, std::chrono::seconds(6)); - - registry.StopReaper(); - EXPECT_TRUE(reaped); - EXPECT_FALSE(registry.IsClientAlive("c1")); -} - -TEST(ClientRegistryTest, ReaperKeepsClientAliveWithHeartbeats) { - ClientRegistry registry(FastExpiryConfig()); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - registry.StartReaper(); - - const auto start = std::chrono::steady_clock::now(); - while (std::chrono::steady_clock::now() - start < std::chrono::seconds(3)) { - EXPECT_EQ(Beat(registry, "c1", Caps(80, 48)), ClientStatus::ALIVE); - std::this_thread::sleep_for(std::chrono::milliseconds(300)); - } - - registry.StopReaper(); - EXPECT_EQ(registry.ClientCount(), 1u); - EXPECT_TRUE(registry.IsClientAlive("c1")); -} - -TEST(ClientRegistryTest, ReaperSelectiveExpiry) { - ClientRegistry registry(FastExpiryConfig()); - EXPECT_TRUE(registry.RegisterClient("c1", "addr-1", Caps(80, 64))); - EXPECT_TRUE(registry.RegisterClient("c2", "addr-2", Caps(80, 64))); - registry.StartReaper(); - - // Keep c1 fed; let c2 go silent. c2 must be reaped while c1 survives. - const bool reached = WaitUntil( - [®istry] { - Beat(registry, "c1", Caps(80, 48)); - return registry.IsClientAlive("c1") && !registry.IsClientAlive("c2"); - }, - std::chrono::seconds(6), std::chrono::milliseconds(200)); - - registry.StopReaper(); - EXPECT_TRUE(reached); - EXPECT_TRUE(registry.IsClientAlive("c1")); - EXPECT_FALSE(registry.IsClientAlive("c2")); -} - -TEST(ClientRegistryTest, StopReaperWhenNeverStarted) { - ClientRegistry registry(ClientRegistryConfig{}); - registry.StopReaper(); // must not hang or crash - SUCCEED(); -} - -TEST(ClientRegistryTest, StartStopReaperMultiple) { - ClientRegistry registry(ClientRegistryConfig{}); - registry.StartReaper(); - registry.StopReaper(); - registry.StartReaper(); - registry.StopReaper(); - SUCCEED(); -} - -TEST(ClientRegistryTest, DestructorStopsRunningReaper) { - ClientRegistry registry(ClientRegistryConfig{}); - registry.StartReaper(); - EXPECT_TRUE(registry.RegisterClient("c1", "addr", Caps(80, 64))); - // Falling out of scope must join the reaper thread cleanly. -} - -} // namespace mori::umbp diff --git a/src/umbp/tests/test_client_registry_external_kv.cpp b/src/umbp/tests/test_client_registry_external_kv.cpp deleted file mode 100644 index e232fe58a..000000000 --- a/src/umbp/tests/test_client_registry_external_kv.cpp +++ /dev/null @@ -1,55 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/master/external_kv_block_index.h" -#include "umbp/distributed/master/global_block_index.h" - -namespace mori::umbp { - -TEST(ClientRegistryExternalKv, UnregisterClientClearsBothIndices) { - GlobalBlockIndex global_index; - ExternalKvBlockIndex external_index; - ClientRegistry registry(ClientRegistryConfig{}, global_index, &external_index); - - ASSERT_TRUE(registry.RegisterClient("node-A", "127.0.0.1:9000", {}, "127.0.0.1:9001")); - ASSERT_EQ(global_index.ApplyEvents("node-A", - {KvEvent{KvEvent::Kind::ADD, "owned", TierType::DRAM, 128}}), - 1u); - ASSERT_EQ(external_index.Register("node-A", {"external"}, TierType::DRAM), 1u); - - registry.UnregisterClient("node-A"); - - EXPECT_TRUE(global_index.Lookup("owned").empty()); - EXPECT_TRUE(external_index.Match({"external"}).empty()); -} - -TEST(ClientRegistryExternalKv, UnregisterWithoutExternalIndexDoesNotCrash) { - GlobalBlockIndex global_index; - ClientRegistry registry(ClientRegistryConfig{}, global_index); - - ASSERT_TRUE(registry.RegisterClient("node-A", "127.0.0.1:9000", {}, "127.0.0.1:9001")); - EXPECT_NO_THROW(registry.UnregisterClient("node-A")); -} - -} // namespace mori::umbp diff --git a/src/umbp/tests/test_external_kv_block_index.cpp b/src/umbp/tests/test_external_kv_block_index.cpp deleted file mode 100644 index 18ee654d5..000000000 --- a/src/umbp/tests/test_external_kv_block_index.cpp +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include -#include -#include - -#include "umbp/distributed/master/external_kv_block_index.h" - -namespace mori::umbp { -namespace { - -const ExternalKvBlockIndex::NodeMatch* FindMatch( - const std::vector& matches, const std::string& node_id) { - for (const auto& match : matches) { - if (match.node_id == node_id) return &match; - } - return nullptr; -} - -std::vector Sorted(std::vector values) { - std::sort(values.begin(), values.end()); - return values; -} - -} // namespace - -TEST(ExternalKvBlockIndex, RegisterIsAdditiveAcrossTiersAndCountsMutations) { - ExternalKvBlockIndex index; - - EXPECT_EQ(index.Register("node-A", {"h1"}, TierType::HBM), 1u); - EXPECT_EQ(index.Register("node-A", {"h1"}, TierType::DRAM), 1u); - EXPECT_EQ(index.Register("node-A", {"h1"}, TierType::DRAM), 0u); - - auto matches = index.Match({"h1"}); - ASSERT_EQ(matches.size(), 1u); - EXPECT_EQ(matches[0].MatchedHashCount(), 1u); - EXPECT_EQ(matches[0].hashes_by_tier.at(TierType::HBM), std::vector({"h1"})); - EXPECT_EQ(matches[0].hashes_by_tier.at(TierType::DRAM), std::vector({"h1"})); - EXPECT_EQ(index.GetKvCount("node-A"), 1u); -} - -TEST(ExternalKvBlockIndex, UnregisterRemovesOnlyRequestedTier) { - ExternalKvBlockIndex index; - ASSERT_EQ(index.Register("node-A", {"h1", "h2"}, TierType::HBM), 2u); - ASSERT_EQ(index.Register("node-A", {"h1"}, TierType::DRAM), 1u); - - EXPECT_EQ(index.Unregister("node-A", {"h1", "missing"}, TierType::HBM), 1u); - EXPECT_EQ(index.Unregister("node-A", {"h1"}, TierType::HBM), 0u); - - auto matches = index.Match({"h1", "h2"}); - ASSERT_EQ(matches.size(), 1u); - const auto& match = matches[0]; - EXPECT_EQ(match.hashes_by_tier.at(TierType::DRAM), std::vector({"h1"})); - EXPECT_EQ(match.hashes_by_tier.at(TierType::HBM), std::vector({"h2"})); - EXPECT_EQ(index.GetKvCount("node-A"), 2u); -} - -TEST(ExternalKvBlockIndex, BulkUnregisterByTierAndNode) { - ExternalKvBlockIndex index; - ASSERT_EQ(index.Register("node-A", {"h1", "h2", "h3"}, TierType::DRAM), 3u); - ASSERT_EQ(index.Register("node-A", {"h1", "h2"}, TierType::SSD), 2u); - ASSERT_EQ(index.Register("node-B", {"h1"}, TierType::SSD), 1u); - - EXPECT_EQ(index.UnregisterByNodeAtTier("node-A", TierType::SSD), 2u); - auto matches = index.Match({"h1", "h2", "h3"}); - ASSERT_EQ(matches.size(), 2u); - const auto* node_a = FindMatch(matches, "node-A"); - ASSERT_NE(node_a, nullptr); - ASSERT_EQ(node_a->hashes_by_tier.size(), 1u); - EXPECT_EQ(Sorted(node_a->hashes_by_tier.at(TierType::DRAM)), - (std::vector{"h1", "h2", "h3"})); - const auto* node_b = FindMatch(matches, "node-B"); - ASSERT_NE(node_b, nullptr); - EXPECT_EQ(node_b->hashes_by_tier.at(TierType::SSD), std::vector({"h1"})); - - EXPECT_EQ(index.UnregisterByNode("node-A"), 3u); - matches = index.Match({"h1", "h2", "h3"}); - ASSERT_EQ(matches.size(), 1u); - EXPECT_EQ(matches[0].node_id, "node-B"); -} - -} // namespace mori::umbp diff --git a/src/umbp/tests/test_external_kv_hit_index.cpp b/src/umbp/tests/test_external_kv_hit_index.cpp deleted file mode 100644 index 20685f0b7..000000000 --- a/src/umbp/tests/test_external_kv_hit_index.cpp +++ /dev/null @@ -1,116 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/master/external_kv_hit_index.h" - -namespace mori::umbp { -namespace { - -std::unordered_map LookupMap(ExternalKvHitIndex& index, - const std::vector& hashes) { - std::vector> entries; - index.Lookup(hashes, &entries); - std::unordered_map out; - for (const auto& [hash, total] : entries) out[hash] = total; - return out; -} - -TEST(ExternalKvHitIndexTest, IncrementAndLookup) { - ExternalKvHitIndex index; - index.IncrementHits({"h1", "h2"}, 100); - - auto counts = LookupMap(index, {"h1", "h2", "missing"}); - ASSERT_EQ(counts.size(), 2); - EXPECT_EQ(counts["h1"], 1); - EXPECT_EQ(counts["h2"], 1); -} - -TEST(ExternalKvHitIndexTest, RepeatedIncrementsAccumulate) { - ExternalKvHitIndex index; - for (int i = 0; i < 10; ++i) index.IncrementHits({"hot"}, 100 + i); - - auto counts = LookupMap(index, {"hot"}); - ASSERT_EQ(counts.size(), 1); - EXPECT_EQ(counts["hot"], 10); -} - -TEST(ExternalKvHitIndexTest, LookupSkipsMissingAndDedupesRequestHashes) { - ExternalKvHitIndex index; - index.IncrementHits({"h1"}, 100); - - std::vector> entries; - index.Lookup({"missing", "h1", "h1", "missing"}, &entries); - ASSERT_EQ(entries.size(), 1); - EXPECT_EQ(entries[0].first, "h1"); - EXPECT_EQ(entries[0].second, 1); -} - -TEST(ExternalKvHitIndexTest, GarbageCollectUsesLastSeenCutoff) { - ExternalKvHitIndex index; - index.IncrementHits({"old"}, 100); - index.IncrementHits({"fresh"}, 200); - - EXPECT_EQ(index.GarbageCollect(150), 1); - EXPECT_EQ(index.Size(), 1); - - auto counts = LookupMap(index, {"old", "fresh"}); - ASSERT_EQ(counts.size(), 1); - EXPECT_EQ(counts["fresh"], 1); -} - -TEST(ExternalKvHitIndexTest, ConcurrentCreationKeepsAllIncrements) { - ExternalKvHitIndex index; - constexpr int kThreads = 32; - constexpr int kIterations = 1000; - - std::atomic start{false}; - std::vector threads; - threads.reserve(kThreads); - for (int t = 0; t < kThreads; ++t) { - threads.emplace_back([&] { - while (!start.load(std::memory_order_acquire)) { - std::this_thread::yield(); - } - for (int i = 0; i < kIterations; ++i) { - index.IncrementHits({"shared"}, static_cast(100 + i)); - } - }); - } - - start.store(true, std::memory_order_release); - for (auto& thread : threads) thread.join(); - - auto counts = LookupMap(index, {"shared"}); - ASSERT_EQ(counts.size(), 1); - EXPECT_EQ(counts["shared"], static_cast(kThreads * kIterations)); -} - -} // namespace -} // namespace mori::umbp diff --git a/src/umbp/tests/test_global_block_index_events.cpp b/src/umbp/tests/test_global_block_index_events.cpp deleted file mode 100644 index b69e7464d..000000000 --- a/src/umbp/tests/test_global_block_index_events.cpp +++ /dev/null @@ -1,489 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. -#include - -#include -#include -#include -#include - -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/master/global_block_index.h" -#include "umbp/distributed/types.h" - -namespace mori::umbp { - -namespace { - -KvEvent Add(std::string key, TierType tier, uint64_t size) { - return KvEvent{KvEvent::Kind::ADD, std::move(key), tier, size}; -} - -KvEvent Remove(std::string key, TierType tier) { - return KvEvent{KvEvent::Kind::REMOVE, std::move(key), tier, 0}; -} - -EventBundle Bundle(uint64_t seq, std::vector events) { - return EventBundle{seq, std::move(events)}; -} - -bool HasLocation(const std::vector& locs, const std::string& node, TierType tier, - uint64_t size) { - for (const auto& l : locs) { - if (l.node_id == node && l.tier == tier && l.size == size) return true; - } - return false; -} - -} // namespace - -// ---- ApplyEvents: ADD/REMOVE round-trip ------------------------------------ - -TEST(GlobalBlockIndexEvents, ApplyAddInsertsLocation) { - GlobalBlockIndex idx; - ASSERT_EQ(idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 1024)}), 1u); - auto locs = idx.Lookup("k1"); - ASSERT_EQ(locs.size(), 1u); - EXPECT_EQ(locs[0].node_id, "node-A"); - EXPECT_EQ(locs[0].tier, TierType::DRAM); - EXPECT_EQ(locs[0].size, 1024u); -} - -// Duplicate ADD keeps the first observed size; only REMOVE retires it. -TEST(GlobalBlockIndexEvents, ApplyAddSameNodeTierKeepsExistingSize) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 1024)}); - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 2048)}); - auto locs = idx.Lookup("k"); - ASSERT_EQ(locs.size(), 1u); - EXPECT_EQ(locs[0].size, 1024u); -} - -TEST(GlobalBlockIndexEvents, MultipleNodesCoexistForSameKey) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Add("k", TierType::DRAM, 200)}); - idx.ApplyEvents("node-A", {Add("k", TierType::HBM, 300)}); // different tier on A - auto locs = idx.Lookup("k"); - EXPECT_EQ(locs.size(), 3u); - EXPECT_TRUE(HasLocation(locs, "node-A", TierType::DRAM, 100)); - EXPECT_TRUE(HasLocation(locs, "node-B", TierType::DRAM, 200)); - EXPECT_TRUE(HasLocation(locs, "node-A", TierType::HBM, 300)); -} - -TEST(GlobalBlockIndexEvents, RemoveErasesMatchingLocationOnly) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Add("k", TierType::DRAM, 200)}); - - idx.ApplyEvents("node-A", {Remove("k", TierType::DRAM)}); - auto locs = idx.Lookup("k"); - ASSERT_EQ(locs.size(), 1u); - EXPECT_EQ(locs[0].node_id, "node-B"); -} - -TEST(GlobalBlockIndexEvents, RemoveLastLocationErasesEntry) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-A", {Remove("k", TierType::DRAM)}); - EXPECT_TRUE(idx.Lookup("k").empty()); - EXPECT_FALSE(idx.GetMetrics("k").has_value()); -} - -TEST(GlobalBlockIndexEvents, RemoveUnknownIsNoop) { - GlobalBlockIndex idx; - EXPECT_EQ(idx.ApplyEvents("ghost", {Remove("ghost-key", TierType::DRAM)}), 0u); -} - -// A key mirrored on both DRAM and SSD of one node: a DRAM eviction -// (REMOVE DRAM) must drop only the DRAM bucket and leave the SSD location -// readable. This is the additive-index invariant the SSD tier relies -// on (DRAM evict never touches the SSD copy); no master code is exercised here -// beyond ApplyEvents. -TEST(GlobalBlockIndexEvents, RemoveDramKeepsSsdBucket) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100), Add("k", TierType::SSD, 100)}); - ASSERT_TRUE(HasLocation(idx.Lookup("k"), "node-A", TierType::DRAM, 100)); - ASSERT_TRUE(HasLocation(idx.Lookup("k"), "node-A", TierType::SSD, 100)); - - idx.ApplyEvents("node-A", {Remove("k", TierType::DRAM)}); - - auto locs = idx.Lookup("k"); - EXPECT_FALSE(HasLocation(locs, "node-A", TierType::DRAM, 100)); // DRAM bucket gone - EXPECT_TRUE(HasLocation(locs, "node-A", TierType::SSD, 100)); // SSD bucket retained -} - -// ---- ReplaceNodeLocations: full-sync recovery ------------------------------ - -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsClearsThenInserts) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 100), Add("k2", TierType::DRAM, 200)}); - idx.ApplyEvents("node-B", {Add("k1", TierType::DRAM, 999)}); // shared key, different node - - // Full-sync from node-A: k1 stays (different size), k2 is gone, new k3 appears. - idx.ReplaceNodeLocations("node-A", - {Add("k1", TierType::DRAM, 150), Add("k3", TierType::DRAM, 300)}); - - auto k1 = idx.Lookup("k1"); - EXPECT_TRUE(HasLocation(k1, "node-A", TierType::DRAM, 150)); - EXPECT_TRUE(HasLocation(k1, "node-B", TierType::DRAM, 999)); // node-B untouched - - EXPECT_TRUE(idx.Lookup("k2").empty()); // dropped — node-A's full-sync didn't include it - - auto k3 = idx.Lookup("k3"); - EXPECT_TRUE(HasLocation(k3, "node-A", TierType::DRAM, 300)); -} - -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsEmptyClearsAllForNode) { - // Used by ClientRegistry::UnregisterClient and the reaper to drop a - // dead node's index entries. - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 1), Add("k2", TierType::HBM, 2)}); - idx.ApplyEvents("node-B", {Add("k1", TierType::DRAM, 3)}); - - idx.ReplaceNodeLocations("node-A", {}); - EXPECT_EQ(idx.Lookup("k1").size(), 1u); // node-B still owns k1 - EXPECT_TRUE(idx.Lookup("k2").empty()); // node-A's only HBM location is gone -} - -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsIgnoresRemoveEntries) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 100)}); - // Snapshot full-sync conventionally carries only ADDs; sneaking - // a REMOVE in is silently skipped (the snapshot is the truth). - idx.ReplaceNodeLocations("node-A", - {Add("k2", TierType::DRAM, 200), Remove("k3", TierType::DRAM)}); - EXPECT_TRUE(idx.Lookup("k1").empty()); - EXPECT_FALSE(idx.Lookup("k2").empty()); - EXPECT_TRUE(idx.Lookup("k3").empty()); -} - -// ---- Reverse-index (node_to_keys_) invariants ------------------------------ - -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsAfterMultiTierRemoveKeepsKeyClean) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100), Add("k", TierType::HBM, 200)}); - idx.ApplyEvents("node-B", {Add("k", TierType::DRAM, 300)}); - - // A still owns (k, HBM): reverse index must keep k. - idx.ApplyEvents("node-A", {Remove("k", TierType::DRAM)}); - auto mid = idx.Lookup("k"); - ASSERT_EQ(mid.size(), 2u); - EXPECT_TRUE(HasLocation(mid, "node-A", TierType::HBM, 200)); - EXPECT_TRUE(HasLocation(mid, "node-B", TierType::DRAM, 300)); - - idx.ReplaceNodeLocations("node-A", {}); - auto after = idx.Lookup("k"); - ASSERT_EQ(after.size(), 1u); - EXPECT_EQ(after[0].node_id, "node-B"); - EXPECT_EQ(after[0].tier, TierType::DRAM); - EXPECT_EQ(after[0].size, 300u); -} - -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsLeavesOtherNodesIntact) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 1), Add("k2", TierType::DRAM, 2)}); - idx.ApplyEvents("node-B", {Add("k1", TierType::DRAM, 10), Add("k3", TierType::HBM, 30)}); - idx.ApplyEvents("node-C", {Add("k2", TierType::HBM, 200), Add("k4", TierType::DRAM, 400)}); - - idx.ReplaceNodeLocations("node-A", {Add("k_new", TierType::DRAM, 999)}); - - auto k1 = idx.Lookup("k1"); - ASSERT_EQ(k1.size(), 1u); - EXPECT_EQ(k1[0].node_id, "node-B"); - EXPECT_EQ(k1[0].size, 10u); - - auto k2 = idx.Lookup("k2"); - ASSERT_EQ(k2.size(), 1u); - EXPECT_EQ(k2[0].node_id, "node-C"); - EXPECT_EQ(k2[0].tier, TierType::HBM); - - EXPECT_TRUE(HasLocation(idx.Lookup("k_new"), "node-A", TierType::DRAM, 999)); - - auto k3 = idx.Lookup("k3"); - ASSERT_EQ(k3.size(), 1u); - EXPECT_EQ(k3[0].node_id, "node-B"); - EXPECT_EQ(k3[0].size, 30u); - - auto k4 = idx.Lookup("k4"); - ASSERT_EQ(k4.size(), 1u); - EXPECT_EQ(k4[0].node_id, "node-C"); - EXPECT_EQ(k4[0].size, 400u); -} - -// 2nd sync must see reverse index repopulated by 1st sync's replay. -TEST(GlobalBlockIndexEvents, ReplaceNodeLocationsTwiceRotatesKeys) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k_old", TierType::DRAM, 1)}); - - idx.ReplaceNodeLocations("node-A", - {Add("k_mid_a", TierType::DRAM, 2), Add("k_mid_b", TierType::HBM, 3)}); - EXPECT_TRUE(idx.Lookup("k_old").empty()); - EXPECT_FALSE(idx.Lookup("k_mid_a").empty()); - EXPECT_FALSE(idx.Lookup("k_mid_b").empty()); - - idx.ReplaceNodeLocations("node-A", {Add("k_final", TierType::DRAM, 4)}); - EXPECT_TRUE(idx.Lookup("k_mid_a").empty()); - EXPECT_TRUE(idx.Lookup("k_mid_b").empty()); - auto final_locs = idx.Lookup("k_final"); - ASSERT_EQ(final_locs.size(), 1u); - EXPECT_EQ(final_locs[0].node_id, "node-A"); - EXPECT_EQ(final_locs[0].size, 4u); -} - -// Reverse-index insert must run even when inserted==false. -TEST(GlobalBlockIndexEvents, DuplicateAddKeepsReverseConsistent) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("dup", TierType::DRAM, 1024)}); - idx.ApplyEvents("node-A", {Add("dup", TierType::DRAM, 2048)}); - idx.ApplyEvents("node-A", {Add("dup", TierType::DRAM, 4096)}); - - auto locs = idx.Lookup("dup"); - ASSERT_EQ(locs.size(), 1u); - - idx.ReplaceNodeLocations("node-A", {}); - EXPECT_TRUE(idx.Lookup("dup").empty()); -} - -// No-op REMOVE must leave node_to_keys_ untouched on both sides. -TEST(GlobalBlockIndexEvents, RemoveNonMatchingTierLeavesReverseUntouched) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - - idx.ApplyEvents("node-A", {Remove("k", TierType::HBM)}); - auto mid = idx.Lookup("k"); - ASSERT_EQ(mid.size(), 1u); - EXPECT_EQ(mid[0].tier, TierType::DRAM); - - idx.ReplaceNodeLocations("node-A", {}); - EXPECT_TRUE(idx.Lookup("k").empty()); - - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Remove("k", TierType::DRAM)}); - idx.ReplaceNodeLocations("node-B", {}); - auto after = idx.Lookup("k"); - ASSERT_EQ(after.size(), 1u); - EXPECT_EQ(after[0].node_id, "node-A"); -} - -// ---- ClientRegistry::Heartbeat applies events end-to-end -------------------- - -TEST(ClientRegistryHeartbeat, AppliesEventsAndAdvancesSeq) { - GlobalBlockIndex idx; - ClientRegistryConfig cfg; - ClientRegistry reg(cfg, idx); - - ASSERT_TRUE(reg.RegisterClient("node-A", "10.0.0.1:1", /*caps=*/{}, "10.0.0.1:2", {})); - - uint64_t acked = 0; - bool need_full = false; - auto status = reg.Heartbeat("node-A", /*caps=*/{}, {Bundle(1, {Add("k", TierType::DRAM, 42)})}, - /*is_full_sync=*/false, /*delta_seq_baseline=*/0, &acked, &need_full); - EXPECT_EQ(status, ClientStatus::ALIVE); - EXPECT_EQ(acked, 1u); - EXPECT_FALSE(need_full); - EXPECT_FALSE(idx.Lookup("k").empty()); -} - -TEST(ClientRegistryHeartbeat, SeqGapTriggersFullSyncRequest) { - GlobalBlockIndex idx; - ClientRegistryConfig cfg; - ClientRegistry reg(cfg, idx); - reg.RegisterClient("node-A", "10.0.0.1:1", {}, "10.0.0.1:2", {}); - - uint64_t acked = 0; - bool need_full = false; - // First heartbeat seq=1 — applied normally. - reg.Heartbeat("node-A", {}, {Bundle(1, {Add("k1", TierType::DRAM, 1)})}, - /*is_full_sync=*/false, 0, &acked, &need_full); - ASSERT_FALSE(need_full); - ASSERT_EQ(acked, 1u); - - // Second heartbeat skips seq=2: master detects the gap. - reg.Heartbeat("node-A", {}, {Bundle(3, {Add("k2", TierType::DRAM, 2)})}, - /*is_full_sync=*/false, 0, &acked, &need_full); - EXPECT_TRUE(need_full); - EXPECT_EQ(acked, 1u); // unchanged — no events applied from this batch - - // k2 is NOT in the index because the gap-batch was rejected. - EXPECT_TRUE(idx.Lookup("k2").empty()); - EXPECT_FALSE(idx.Lookup("k1").empty()); -} - -TEST(ClientRegistryHeartbeat, FullSyncReplacesNodeLocations) { - GlobalBlockIndex idx; - ClientRegistryConfig cfg; - ClientRegistry reg(cfg, idx); - reg.RegisterClient("node-A", "10.0.0.1:1", {}, "10.0.0.1:2", {}); - - uint64_t acked = 0; - bool need_full = false; - reg.Heartbeat("node-A", {}, - {Bundle(1, {Add("k1", TierType::DRAM, 1), Add("k2", TierType::DRAM, 2)})}, - /*is_full_sync=*/false, 0, &acked, &need_full); - ASSERT_FALSE(idx.Lookup("k1").empty()); - ASSERT_FALSE(idx.Lookup("k2").empty()); - - // Full-sync: only k1 + k3 should remain for node-A. - reg.Heartbeat("node-A", {}, - {Bundle(2, {Add("k1", TierType::DRAM, 10), Add("k3", TierType::DRAM, 30)})}, - /*is_full_sync=*/true, /*delta_seq_baseline=*/2, &acked, &need_full); - EXPECT_EQ(acked, 2u); - EXPECT_FALSE(need_full); - - auto k1 = idx.Lookup("k1"); - ASSERT_EQ(k1.size(), 1u); - EXPECT_EQ(k1[0].size, 10u); // updated via full-sync - EXPECT_TRUE(idx.Lookup("k2").empty()); - EXPECT_FALSE(idx.Lookup("k3").empty()); -} - -TEST(ClientRegistryHeartbeat, UnregisterClearsNodeFromIndex) { - GlobalBlockIndex idx; - ClientRegistryConfig cfg; - ClientRegistry reg(cfg, idx); - reg.RegisterClient("node-A", "10.0.0.1:1", {}, "10.0.0.1:2", {}); - - uint64_t acked = 0; - bool need_full = false; - reg.Heartbeat("node-A", {}, {Bundle(1, {Add("k1", TierType::DRAM, 1)})}, - /*is_full_sync=*/false, 0, &acked, &need_full); - ASSERT_FALSE(idx.Lookup("k1").empty()); - - reg.UnregisterClient("node-A"); - EXPECT_TRUE(idx.Lookup("k1").empty()); - EXPECT_FALSE(reg.IsClientAlive("node-A")); -} - -// ---- FindEvictionCandidates -------------------------------------------------- - -TEST(GlobalBlockIndexEvents, FindEvictionCandidatesFiltersByOverloadedNodeTier) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 100), Add("k2", TierType::HBM, 200)}); - idx.ApplyEvents("node-B", {Add("k1", TierType::DRAM, 100)}); - - std::set overloaded = { - {"node-A", TierType::DRAM}, - }; - auto candidates = idx.FindEvictionCandidates(overloaded); - // Only node-A's DRAM location of k1 is a candidate. - ASSERT_EQ(candidates.size(), 1u); - EXPECT_EQ(candidates[0].key, "k1"); - EXPECT_EQ(candidates[0].location.node_id, "node-A"); - EXPECT_EQ(candidates[0].location.tier, TierType::DRAM); -} - -// ---- BatchLookupForRouteGet ------------------------------------------------ - -TEST(GlobalBlockIndexEvents, BatchLookupForRouteGetEmptyInputReturnsEmpty) { - GlobalBlockIndex idx; - EXPECT_TRUE(idx.BatchLookupForRouteGet({}, {}, std::chrono::seconds{1}).empty()); -} - -TEST(GlobalBlockIndexEvents, BatchLookupForRouteGetMixedHitsAndMisses) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k1", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Add("k1", TierType::DRAM, 200), Add("k2", TierType::HBM, 300)}); - - auto ref_k1 = idx.Lookup("k1"); - auto ref_k2 = idx.Lookup("k2"); - auto before_k1 = idx.GetMetrics("k1"); - auto before_k2 = idx.GetMetrics("k2"); - ASSERT_TRUE(before_k1.has_value()); - ASSERT_TRUE(before_k2.has_value()); - - auto results = idx.BatchLookupForRouteGet({"k1", "ghost", "k2"}, {}, std::chrono::seconds{10}); - ASSERT_EQ(results.size(), 3u); - EXPECT_EQ(results[0], ref_k1); - EXPECT_TRUE(results[1].empty()); - EXPECT_EQ(results[2], ref_k2); - - auto after_k1 = idx.GetMetrics("k1"); - auto after_k2 = idx.GetMetrics("k2"); - ASSERT_TRUE(after_k1.has_value()); - ASSERT_TRUE(after_k2.has_value()); - EXPECT_EQ(after_k1->access_count, before_k1->access_count + 1); - EXPECT_EQ(after_k2->access_count, before_k2->access_count + 1); - EXPECT_FALSE(idx.GetMetrics("ghost").has_value()); -} - -TEST(GlobalBlockIndexEvents, BatchLookupForRouteGetGrantsLeaseForHitsOnly) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("hit", TierType::DRAM, 100), Add("other", TierType::DRAM, 200)}); - - std::set overloaded{{"node-A", TierType::DRAM}}; - ASSERT_EQ(idx.FindEvictionCandidates(overloaded).size(), 2u); - - idx.BatchLookupForRouteGet({"hit", "ghost"}, {}, std::chrono::seconds{10}); - - auto candidates = idx.FindEvictionCandidates(overloaded); - ASSERT_EQ(candidates.size(), 1u); - EXPECT_EQ(candidates[0].key, "other"); -} - -// All replicas excluded -> slot empty, access_count NOT bumped, -// lease NOT granted. A key whose every replica is unreachable must -// not pollute LRU or block eviction. -TEST(GlobalBlockIndexEvents, BatchLookupForRouteGetSkipsSideEffectsWhenAllReplicasExcluded) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Add("k", TierType::DRAM, 200)}); - - auto before = idx.GetMetrics("k"); - ASSERT_TRUE(before.has_value()); - std::set overloaded{{"node-A", TierType::DRAM}, - {"node-B", TierType::DRAM}}; - ASSERT_EQ(idx.FindEvictionCandidates(overloaded).size(), 2u); - - std::unordered_set excludes{"node-A", "node-B"}; - auto results = idx.BatchLookupForRouteGet({"k"}, excludes, std::chrono::seconds{10}); - ASSERT_EQ(results.size(), 1u); - EXPECT_TRUE(results[0].empty()); - - auto after = idx.GetMetrics("k"); - ASSERT_TRUE(after.has_value()); - EXPECT_EQ(after->access_count, before->access_count); - EXPECT_EQ(idx.FindEvictionCandidates(overloaded).size(), 2u); -} - -// Some replicas excluded but not all -> returned slot has only the -// survivors, access_count IS bumped, lease IS granted. -TEST(GlobalBlockIndexEvents, BatchLookupForRouteGetFiltersAndLeasesWhenSomeReplicasSurvive) { - GlobalBlockIndex idx; - idx.ApplyEvents("node-A", {Add("k", TierType::DRAM, 100)}); - idx.ApplyEvents("node-B", {Add("k", TierType::DRAM, 200)}); - - auto before = idx.GetMetrics("k"); - ASSERT_TRUE(before.has_value()); - - std::unordered_set excludes{"node-A"}; - auto results = idx.BatchLookupForRouteGet({"k"}, excludes, std::chrono::seconds{10}); - ASSERT_EQ(results.size(), 1u); - ASSERT_EQ(results[0].size(), 1u); - EXPECT_EQ(results[0][0].node_id, "node-B"); - - auto after = idx.GetMetrics("k"); - ASSERT_TRUE(after.has_value()); - EXPECT_EQ(after->access_count, before->access_count + 1); -} - -} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/CMakeLists.txt b/tests/cpp/umbp/distributed/CMakeLists.txt index 5beb14654..6e8f3da3e 100644 --- a/tests/cpp/umbp/distributed/CMakeLists.txt +++ b/tests/cpp/umbp/distributed/CMakeLists.txt @@ -37,6 +37,10 @@ target_link_libraries( add_test(NAME umbp_master_client_flush_heartbeat COMMAND test_umbp_master_client_flush_heartbeat) +# Enables gtest_discover_tests() for the metadata-store / external-KV / SSD-tier +# suite migrated below from src/umbp/tests. +include(GoogleTest) + # MasterClient lifecycle: destructor budget, heartbeat thread shutdown, etc. add_executable(test_umbp_master_client_lifecycle test_master_client_lifecycle.cpp) @@ -226,16 +230,124 @@ target_link_libraries( PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} ${_TEST_GRPCPP_LIB} gtest_main) -# GlobalBlockIndex::BatchLookupForRouteGet micro-bench (pure logic; no gRPC / -# RDMA, so umbp_common only). Drives the master-side BatchRouteGet hot path -# directly: hash lookup + exclude filter + access bump + lease grant under a -# single shared_lock. Reports us/call, Mlookups/s, and concurrent read scaling -# via --threads. -add_executable(bench_umbp_global_block_index_route_get - bench_global_block_index_route_get.cpp) -target_link_libraries(bench_umbp_global_block_index_route_get - PRIVATE umbp_common gtest_main) +# RouteGet strategy bench: per-key Select() loop vs the batched BatchSelect() +# the router now calls once per BatchRouteGet. Pure logic (umbp_common only); +# standalone executable, not a CTest target. +add_executable(bench_umbp_route_strategy bench_route_strategy.cpp) +target_link_libraries(bench_umbp_route_strategy PRIVATE umbp_common) + +# =========================================================================== +# Metadata-store / external-KV / SSD-tier suite (migrated from src/umbp/tests). +# These link umbp_common (pure logic) except test_peer_ssd_read_rpc, which +# exercises the gRPC peer service and links umbp_core + protobuf/gRPC. Test +# cases are registered individually via gtest_discover_tests. +# =========================================================================== + +# test_master_metadata_store_interface — Phase 1 compile/instantiation gate for +# IMasterMetadataStore. Includes an isolated self-compile TU (proves the header +# is self-contained) and a GMock mock instantiation/signature-completeness test. +add_executable( + test_master_metadata_store_interface test_master_metadata_store_interface.cpp + master_metadata_store_self_compile.cpp) +target_include_directories(test_master_metadata_store_interface + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_link_libraries(test_master_metadata_store_interface + PRIVATE umbp_common GTest::gmock GTest::gtest_main) +target_compile_features(test_master_metadata_store_interface PRIVATE cxx_std_17) +gtest_discover_tests(test_master_metadata_store_interface) + +# test_in_memory_master_metadata_store — Phase 2 behavioral suite for the +# InMemoryMasterMetadataStore implementation of IMasterMetadataStore (§6a). +add_executable(test_in_memory_master_metadata_store + test_in_memory_master_metadata_store.cpp) +target_link_libraries(test_in_memory_master_metadata_store + PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_in_memory_master_metadata_store PRIVATE cxx_std_17) +gtest_discover_tests(test_in_memory_master_metadata_store) + +# test_peer_dram_allocator +add_executable(test_peer_dram_allocator test_peer_dram_allocator.cpp) +target_link_libraries(test_peer_dram_allocator PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_peer_dram_allocator PRIVATE cxx_std_17) +gtest_discover_tests(test_peer_dram_allocator) + +# test_router_dedup — master-side BatchRoutePut dedup via IMasterMetadataStore +add_executable(test_router_dedup test_router_dedup.cpp) +target_link_libraries(test_router_dedup PRIVATE umbp_common GTest::gtest_main) +target_compile_features(test_router_dedup PRIVATE cxx_std_17) +gtest_discover_tests(test_router_dedup) + +# test_peer_ssd_manager — SSD tier ownership + owned-location source +add_executable(test_peer_ssd_manager test_peer_ssd_manager.cpp) +target_link_libraries(test_peer_ssd_manager PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_peer_ssd_manager PRIVATE cxx_std_17) +gtest_discover_tests(test_peer_ssd_manager) + +# test_peer_ssd_eviction — LRU + watermark eviction + in-flight guard +add_executable(test_peer_ssd_eviction test_peer_ssd_eviction.cpp) +target_link_libraries(test_peer_ssd_eviction PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_peer_ssd_eviction PRIVATE cxx_std_17) +gtest_discover_tests(test_peer_ssd_eviction) + +# test_ssd_copy_pipeline — DramCopyPin + async copy-on-commit pipeline +add_executable(test_ssd_copy_pipeline test_ssd_copy_pipeline.cpp) +target_link_libraries(test_ssd_copy_pipeline PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_ssd_copy_pipeline PRIVATE cxx_std_17) +gtest_discover_tests(test_ssd_copy_pipeline) + +# test_tier_priority_route_get — RouteGet tier-priority strategy +add_executable(test_tier_priority_route_get test_tier_priority_route_get.cpp) +target_link_libraries(test_tier_priority_route_get PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_tier_priority_route_get PRIVATE cxx_std_17) +gtest_discover_tests(test_tier_priority_route_get) +# test_peer_ssd_read_rpc — SSD read RPC over a gRPC loopback. Needs the gRPC +# peer service + generated stubs (umbp_core) and protobuf/gRPC. +add_executable(test_peer_ssd_read_rpc test_peer_ssd_read_rpc.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_peer_ssd_read_rpc PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_peer_ssd_read_rpc PRIVATE umbp_core umbp_common ${_TEST_PROTOBUF_LIB} + ${_TEST_GRPCPP_LIB} GTest::gtest_main) +target_compile_features(test_peer_ssd_read_rpc PRIVATE cxx_std_17) +gtest_discover_tests(test_peer_ssd_read_rpc) + +# test_ssd_read_lease_gating — pure reader-side lease gating decision logic +# (ssd_read_lease.h). Header-only under test (no gRPC / RDMA deps). +add_executable(test_ssd_read_lease_gating test_ssd_read_lease_gating.cpp) +target_link_libraries(test_ssd_read_lease_gating PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_ssd_read_lease_gating PRIVATE cxx_std_17) +gtest_discover_tests(test_ssd_read_lease_gating) + +# test_ssd_reliability — cross-component reliability: owned-source DRAM+SSD +# merge, SSD evict -> REMOVE -> master index convergence, tier-priority over the +# real index, crash-restart discard, and observability counters. +add_executable(test_ssd_reliability test_ssd_reliability.cpp) +target_link_libraries(test_ssd_reliability PRIVATE umbp_common + GTest::gtest_main) +target_compile_features(test_ssd_reliability PRIVATE cxx_std_17) +gtest_discover_tests(test_ssd_reliability) + +# test_batch_resolve_codec — struct-of-arrays BatchResolveKeysResponse wire +# codec (batch_resolve_codec.h): encode/decode round-trip, page flattening, desc +# dedup / omit_descs, malformed-response rejection, and the payload-size win. +# Needs the generated proto (carried by umbp_common) + protobuf. +add_executable(test_batch_resolve_codec test_batch_resolve_codec.cpp) +if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_link_options(test_batch_resolve_codec PRIVATE -Wl,--no-as-needed) +endif() +target_link_libraries( + test_batch_resolve_codec PRIVATE umbp_common ${_TEST_PROTOBUF_LIB} + GTest::gtest_main) +target_compile_features(test_batch_resolve_codec PRIVATE cxx_std_17) +gtest_discover_tests(test_batch_resolve_codec) # Multi-client agent KV-event pressure bench: in-process master + N PoolClients # driving a put -> LLM-gap -> get schedule (--pattern broadcast|rotate), # comparing baseline / compressed-heartbeat / flush kv-event propagation under diff --git a/tests/cpp/umbp/distributed/bench_global_block_index_route_get.cpp b/tests/cpp/umbp/distributed/bench_global_block_index_route_get.cpp deleted file mode 100644 index 404b5bf9f..000000000 --- a/tests/cpp/umbp/distributed/bench_global_block_index_route_get.cpp +++ /dev/null @@ -1,275 +0,0 @@ -// Copyright © Advanced Micro Devices, Inc. All rights reserved. -// -// MIT License -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in all -// copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// GlobalBlockIndex::BatchLookupForRouteGet micro-bench (pure logic; no gRPC / -// RDMA). This is the hot master-side path that BatchRouteGet drives: for a -// batch of keys it does, under a single shared_lock, a hash lookup + -// exclude-node filter + (on a non-empty result) an atomic access-count bump and -// lease grant. Because the lock is shared, multiple caller threads can run the -// path concurrently; --threads measures that read-scaling. -// -// The index is pre-populated via ApplyEvents (the only real mutator) with -// `index-keys` keys, each replicated across `locs-per-key` of `nodes` peers, so -// the per-key location vector and the exclude filter have realistic length. -// -// Usage: -// bench_umbp_global_block_index_route_get -// [--index-keys N] [--batch N] [--nodes N] [--locs-per-key N] -// [--exclude N] [--hit-pct P] [--threads N] [--iters N] [--lease-ms N] -// [--sweep batch=1,8,64,512] -// -// CSV (stdout): -// index_keys,batch,nodes,locs_per_key,exclude,hit_pct,threads,iters, -// wall_us_per_call,Mlookups_s,locs_per_call - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "umbp/distributed/master/global_block_index.h" -#include "umbp/distributed/types.h" - -using mori::umbp::GlobalBlockIndex; -using mori::umbp::KvEvent; -using mori::umbp::Location; -using mori::umbp::TierType; - -namespace { - -struct BenchOpts { - size_t index_keys = 100000; // keys resident in the index - size_t batch = 64; // keys per BatchLookupForRouteGet call - size_t nodes = 8; // distinct peer node_ids holding locations - size_t locs_per_key = 2; // replicas (distinct nodes) per key - size_t exclude = 0; // node_ids in the exclude set - size_t hit_pct = 100; // % of looked-up keys that exist in the index - size_t threads = 1; // concurrent caller threads (shared_lock read scaling) - size_t iters = 2000; // measured calls per thread - size_t lease_ms = 10; // lease_duration handed to the method - std::vector sweep; // batch sweep -}; - -void Usage() { - std::cerr << "Usage: bench_umbp_global_block_index_route_get\n" - << " [--index-keys N] [--batch N] [--nodes N] [--locs-per-key N]\n" - << " [--exclude N] [--hit-pct P] [--threads N] [--iters N] [--lease-ms N]\n" - << " [--sweep batch=1,8,64,512]\n"; -} - -bool ParseArgs(int argc, char** argv, BenchOpts* o) { - for (int i = 1; i < argc; ++i) { - std::string a = argv[i]; - auto next = [&](const char* what) -> const char* { - if (i + 1 >= argc) { - std::cerr << "Missing value for " << what << "\n"; - std::exit(2); - } - return argv[++i]; - }; - auto u64 = [&](const char* what) { return std::strtoull(next(what), nullptr, 10); }; - if (a == "--index-keys") { - o->index_keys = u64("--index-keys"); - } else if (a == "--batch") { - o->batch = u64("--batch"); - } else if (a == "--nodes") { - o->nodes = u64("--nodes"); - } else if (a == "--locs-per-key") { - o->locs_per_key = u64("--locs-per-key"); - } else if (a == "--exclude") { - o->exclude = u64("--exclude"); - } else if (a == "--hit-pct") { - o->hit_pct = u64("--hit-pct"); - } else if (a == "--threads") { - o->threads = u64("--threads"); - } else if (a == "--iters") { - o->iters = u64("--iters"); - } else if (a == "--lease-ms") { - o->lease_ms = u64("--lease-ms"); - } else if (a == "--sweep") { - std::string s = next("--sweep"); - auto eq = s.find('='); - if (eq == std::string::npos) { - std::cerr << "--sweep expects 'batch=N1,N2,...'\n"; - return false; - } - std::stringstream ss(s.substr(eq + 1)); - std::string tok; - while (std::getline(ss, tok, ',')) { - if (!tok.empty()) o->sweep.push_back(std::strtoull(tok.c_str(), nullptr, 10)); - } - } else if (a == "-h" || a == "--help") { - Usage(); - std::exit(0); - } else { - std::cerr << "Unknown arg: " << a << "\n"; - Usage(); - return false; - } - } - if (o->locs_per_key < 1) o->locs_per_key = 1; - if (o->nodes < o->locs_per_key) o->nodes = o->locs_per_key; - if (o->threads < 1) o->threads = 1; - if (o->hit_pct > 100) o->hit_pct = 100; - return true; -} - -std::string KeyName(size_t i) { return "k" + std::to_string(i); } -std::string NodeName(size_t n) { return "node-" + std::to_string(n); } - -// Populate the index with `index_keys` keys. Key i is owned by nodes -// (i + 0 .. i + locs_per_key-1) mod nodes, so replicas are spread evenly and -// every node carries a comparable share. We batch all ADDs for a given node -// into one ApplyEvents call (ApplyEvents groups by node_id). -void Populate(GlobalBlockIndex* idx, const BenchOpts& o) { - std::vector> per_node(o.nodes); - for (size_t i = 0; i < o.index_keys; ++i) { - for (size_t r = 0; r < o.locs_per_key; ++r) { - size_t node = (i + r) % o.nodes; - KvEvent ev; - ev.kind = KvEvent::Kind::ADD; - ev.key = KeyName(i); - ev.tier = TierType::DRAM; - ev.size = 4096; - per_node[node].push_back(std::move(ev)); - } - } - for (size_t n = 0; n < o.nodes; ++n) { - idx->ApplyEvents(NodeName(n), per_node[n]); - } -} - -// Pre-build the per-iteration key batches so the measured loop touches only the -// method under test (no string formatting or RNG in the hot path). `hit_pct` -// of the keys index into the populated range; the rest are guaranteed misses. -std::vector> BuildBatches(const BenchOpts& o, uint64_t seed) { - std::mt19937_64 rng(seed); - std::uniform_int_distribution hit_dist(0, o.index_keys ? o.index_keys - 1 : 0); - std::uniform_int_distribution pct(1, 100); - std::vector> batches(o.iters); - for (size_t it = 0; it < o.iters; ++it) { - auto& b = batches[it]; - b.reserve(o.batch); - for (size_t j = 0; j < o.batch; ++j) { - if (pct(rng) <= o.hit_pct && o.index_keys > 0) { - b.push_back(KeyName(hit_dist(rng))); - } else { - b.push_back("miss-" + std::to_string(it) + "-" + std::to_string(j)); - } - } - } - return batches; -} - -std::unordered_set BuildExclude(const BenchOpts& o) { - std::unordered_set ex; - for (size_t n = 0; n < o.exclude && n < o.nodes; ++n) ex.insert(NodeName(n)); - return ex; -} - -void RunScenario(const BenchOpts& base, size_t batch_override) { - BenchOpts o = base; - o.batch = batch_override > 0 ? batch_override : o.batch; - - GlobalBlockIndex idx; - Populate(&idx, o); - const auto exclude = BuildExclude(o); - const auto lease = std::chrono::milliseconds(o.lease_ms); - - // Per-thread inputs (distinct seeds so threads don't all hammer one bucket). - std::vector>> thread_batches(o.threads); - for (size_t t = 0; t < o.threads; ++t) { - thread_batches[t] = BuildBatches(o, 0x9E3779B97F4A7C15ULL ^ (t + 1)); - } - - // Warm-up (not measured): one pass per thread's first batch, single-threaded. - for (size_t t = 0; t < o.threads; ++t) { - (void)idx.BatchLookupForRouteGet(thread_batches[t][0], exclude, lease); - } - - std::atomic total_locs{0}; - std::atomic max_wall_ms{0.0}; - - auto worker = [&](size_t t) { - const auto& batches = thread_batches[t]; - uint64_t locs = 0; - auto t0 = std::chrono::steady_clock::now(); - for (size_t it = 0; it < o.iters; ++it) { - auto out = idx.BatchLookupForRouteGet(batches[it], exclude, lease); - for (const auto& v : out) locs += v.size(); - } - auto t1 = std::chrono::steady_clock::now(); - total_locs.fetch_add(locs, std::memory_order_relaxed); - const double ms = std::chrono::duration(t1 - t0).count(); - // Wall clock is the slowest thread (calls run concurrently under shared_lock). - double prev = max_wall_ms.load(std::memory_order_relaxed); - while (ms > prev && !max_wall_ms.compare_exchange_weak(prev, ms)) { - } - }; - - std::vector pool; - pool.reserve(o.threads); - auto wall0 = std::chrono::steady_clock::now(); - for (size_t t = 0; t < o.threads; ++t) pool.emplace_back(worker, t); - for (auto& th : pool) th.join(); - auto wall1 = std::chrono::steady_clock::now(); - - const double wall_ms = std::chrono::duration(wall1 - wall0).count(); - const uint64_t total_calls = static_cast(o.iters) * o.threads; - const uint64_t total_keys = total_calls * o.batch; - const double us_per_call = total_calls > 0 ? (wall_ms * 1000.0) / total_calls : 0.0; - const double mlookups_s = wall_ms > 0 ? (total_keys / 1e6) / (wall_ms / 1000.0) : 0.0; - const double locs_per_call = - total_calls > 0 ? static_cast(total_locs.load()) / total_calls : 0.0; - - std::printf("%zu,%zu,%zu,%zu,%zu,%zu,%zu,%zu,%.4f,%.2f,%.2f\n", o.index_keys, o.batch, o.nodes, - o.locs_per_key, o.exclude, o.hit_pct, o.threads, o.iters, us_per_call, mlookups_s, - locs_per_call); - std::fflush(stdout); -} - -} // namespace - -int main(int argc, char** argv) { - BenchOpts opts; - if (!ParseArgs(argc, argv, &opts)) return 2; - - std::printf( - "index_keys,batch,nodes,locs_per_key,exclude,hit_pct,threads,iters," - "wall_us_per_call,Mlookups_s,locs_per_call\n"); - std::fflush(stdout); - - if (!opts.sweep.empty()) { - for (size_t b : opts.sweep) RunScenario(opts, b); - } else { - RunScenario(opts, 0); - } - return 0; -} diff --git a/tests/cpp/umbp/distributed/bench_route_strategy.cpp b/tests/cpp/umbp/distributed/bench_route_strategy.cpp new file mode 100644 index 000000000..3fb1640d4 --- /dev/null +++ b/tests/cpp/umbp/distributed/bench_route_strategy.cpp @@ -0,0 +1,217 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// RouteGet strategy micro-bench: per-key Select() loop (the old router hot path) +// vs the batched BatchSelect() the router now calls once per batch. Pure CPU, +// no master/RDMA — isolates the routing-strategy cost the master pays under +// BatchRouteGet. +// +// The overhead the batch path removes is one virtual dispatch per batch instead +// of one per key. Logging is suppressed by the default module level (ERROR); +// set MORI_UMBP_LOG_LEVEL=error explicitly if your environment overrides it. +// +// (The RoutePut path is intentionally not benched here: the put strategy exposes +// only the batched SelectBatch — there is no per-key variant left to compare +// against. See bench_route_put_strategy.cpp for put-side benchmarks.) +// +// Usage: +// bench_umbp_route_strategy [--keys N] [--iters N] [--replicas R] +// [--get-strategy random|tier-priority] +// [--sweep keys=1,4,16,64,256,1024] +// +// CSV (stdout): +// op,keys,iters,total_ms,ns_per_key +// where op is Get / BatchGet (per-key Select() loop vs BatchSelect()). +// total_ms is the duration of ONE iteration (a single batch call over all +// `keys`), averaged across `iters` repetitions; ns_per_key is that divided by +// `keys`. + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/routing/route_get_strategy.h" +#include "umbp/distributed/types.h" + +using mori::umbp::Location; +using mori::umbp::RandomRouteGetStrategy; +using mori::umbp::RouteGetStrategy; +using mori::umbp::TierPriorityRouteGetStrategy; +using mori::umbp::TierType; + +namespace { + +struct BenchOpts { + size_t keys = 256; + size_t iters = 2000; + size_t replicas = 3; // candidate replicas per key on the Get path + std::string get_strategy = "random"; // random | tier-priority + std::vector sweep; +}; + +void Usage() { + std::cerr << "Usage: bench_umbp_route_strategy [--keys N] [--iters N] [--replicas R]\n" + << " [--get-strategy random|tier-priority]\n" + << " [--sweep keys=1,4,16,64,256,1024]\n"; +} + +// Resolve the --get-strategy name to a concrete RouteGetStrategy. +std::unique_ptr MakeGetStrategy(const std::string& name) { + if (name == "random") return std::make_unique(); + if (name == "tier-priority") return std::make_unique(); + std::cerr << "Unknown --get-strategy: " << name << " (expected random|tier-priority)\n"; + std::exit(2); +} + +bool ParseArgs(int argc, char** argv, BenchOpts* o) { + for (int i = 1; i < argc; ++i) { + std::string a = argv[i]; + auto next = [&](const char* what) -> const char* { + if (i + 1 >= argc) { + std::cerr << "Missing value for " << what << "\n"; + std::exit(2); + } + return argv[++i]; + }; + if (a == "--keys") { + o->keys = std::strtoull(next("--keys"), nullptr, 10); + } else if (a == "--iters") { + o->iters = std::strtoull(next("--iters"), nullptr, 10); + } else if (a == "--replicas") { + o->replicas = std::strtoull(next("--replicas"), nullptr, 10); + } else if (a == "--get-strategy") { + o->get_strategy = next("--get-strategy"); + } else if (a == "--sweep") { + std::string s = next("--sweep"); + auto eq = s.find('='); + if (eq == std::string::npos) { + std::cerr << "--sweep expects 'keys=N1,N2,...'\n"; + return false; + } + std::stringstream ss(s.substr(eq + 1)); + std::string tok; + while (std::getline(ss, tok, ',')) { + if (!tok.empty()) o->sweep.push_back(std::strtoull(tok.c_str(), nullptr, 10)); + } + } else if (a == "-h" || a == "--help") { + Usage(); + std::exit(0); + } else { + std::cerr << "Unknown arg: " << a << "\n"; + Usage(); + return false; + } + } + return true; +} + +// --- Get-path workload ----------------------------------------------------- +// One candidate list per key; replicas spread across tiers/nodes so the +// TierPriority strategy does real work (rank scan + within-tier random pick). +std::vector> BuildGetWorkload(const BenchOpts& o) { + static const TierType kTiers[] = {TierType::HBM, TierType::DRAM, TierType::SSD}; + std::vector> per_key(o.keys); + for (size_t k = 0; k < o.keys; ++k) { + auto& locs = per_key[k]; + locs.reserve(o.replicas); + for (size_t r = 0; r < o.replicas; ++r) { + Location loc; + loc.node_id = "node-" + std::to_string((k + r) % 16); + loc.size = 4096; + loc.tier = kTiers[(k + r) % 3]; + locs.push_back(loc); + } + } + return per_key; +} + +// Returns total wall time (ns) over all `iters` measured calls (warm-up excluded). +template +double TotalNs(size_t iters, Fn&& fn) { + fn(); // warm-up (not measured) + auto t0 = std::chrono::steady_clock::now(); + for (size_t it = 0; it < iters; ++it) fn(); + auto t1 = std::chrono::steady_clock::now(); + return std::chrono::duration(t1 - t0).count(); +} + +void Emit(const char* op, size_t keys, size_t iters, double total_ns) { + // Per-iteration duration: one batch call over `keys`, averaged across iters. + const double ns_per_batch = iters ? total_ns / static_cast(iters) : 0.0; + const double ns_per_key = keys ? ns_per_batch / static_cast(keys) : 0.0; + std::printf("%s,%zu,%zu,%.5f,%.2f\n", op, keys, iters, ns_per_batch / 1e6, ns_per_key); + std::fflush(stdout); +} + +void RunGet(const BenchOpts& base, size_t keys) { + BenchOpts o = base; + o.keys = keys; + auto per_key = BuildGetWorkload(o); + auto strat_ptr = MakeGetStrategy(o.get_strategy); + RouteGetStrategy& strat = *strat_ptr; + const std::string node_id = "requester"; + + // Sink to keep the optimizer from eliding the work. + volatile uint64_t sink = 0; + + double get_ns = TotalNs(o.iters, [&] { + for (size_t k = 0; k < per_key.size(); ++k) { + if (per_key[k].empty()) continue; + Location sel = strat.Select(per_key[k], node_id); + sink += sel.size + sel.node_id.size(); + } + }); + + double batch_get_ns = TotalNs(o.iters, [&] { + auto out = strat.BatchSelect(per_key, node_id); + for (const auto& sel : out) sink += sel.size + sel.node_id.size(); + }); + + (void)sink; + Emit("Get", o.keys, o.iters, get_ns); + Emit("BatchGet", o.keys, o.iters, batch_get_ns); +} + +} // namespace + +int main(int argc, char** argv) { + BenchOpts opts; + if (!ParseArgs(argc, argv, &opts)) return 2; + + std::fprintf(stderr, "[config] replicas/key=%zu iters=%zu get-strategy=%s\n", opts.replicas, + opts.iters, opts.get_strategy.c_str()); + std::printf("op,keys,iters,total_ms,ns_per_key\n"); + std::fflush(stdout); + + if (!opts.sweep.empty()) { + for (size_t k : opts.sweep) RunGet(opts, k); + } else { + RunGet(opts, opts.keys); + } + return 0; +} diff --git a/tests/cpp/umbp/distributed/master_metadata_store_self_compile.cpp b/tests/cpp/umbp/distributed/master_metadata_store_self_compile.cpp new file mode 100644 index 000000000..e8e016ce8 --- /dev/null +++ b/tests/cpp/umbp/distributed/master_metadata_store_self_compile.cpp @@ -0,0 +1,27 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Phase 1 "header self-compiles" gate: this translation unit includes ONLY the +// interface header (which in turn includes types.h). If it compiles, the header +// is self-contained — no missing includes or forward declarations. Deliberately +// has no other includes and no symbols of its own. +#include "umbp/distributed/master/master_metadata_store.h" diff --git a/tests/cpp/umbp/distributed/mock_master_metadata_store.h b/tests/cpp/umbp/distributed/mock_master_metadata_store.h new file mode 100644 index 000000000..811e5919a --- /dev/null +++ b/tests/cpp/umbp/distributed/mock_master_metadata_store.h @@ -0,0 +1,126 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// GMock mock for IMasterMetadataStore. +// +// Phase 1 use: instantiation gate. If this type compiles and instantiates, +// every pure-virtual on the interface is overridden with a well-typed +// signature — proving the contract has no orphaned/ill-typed methods. +// +// Reused in Phase 3 (consumer-integration) to assert that each rewired +// consumer (Router / EvictionManager / UMBPMasterServiceImpl handlers) calls +// the right store method with correctly-translated arguments. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/master_metadata_store.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +class MockMasterMetadataStore : public IMasterMetadataStore { + public: + // Aliases for types whose commas would otherwise break the MOCK_METHOD macro + // parser (it splits the argument list on top-level commas). + using CapsMap = std::map; + using LocationBatch = std::vector>; + + // --- Cross-store writes --- + MOCK_METHOD(bool, RegisterClient, + (const ClientRegistration& registration, std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration stale_after), + (override)); + MOCK_METHOD(void, UnregisterClient, (const std::string& node_id), (override)); + MOCK_METHOD(HeartbeatResult, ApplyHeartbeat, + (const std::string& node_id, uint64_t seq, std::chrono::system_clock::time_point now, + const CapsMap& caps, (const std::vector&)events, bool is_full_sync), + (override)); + MOCK_METHOD(std::vector, ExpireStaleClients, + (std::chrono::system_clock::time_point cutoff), (override)); + + // --- External-KV writes --- + MOCK_METHOD(bool, RegisterExternalKvIfAlive, + (const std::string& node_id, (const std::vector&)hashes, TierType tier), + (override)); + MOCK_METHOD(void, UnregisterExternalKv, + (const std::string& node_id, (const std::vector&)hashes, TierType tier), + (override)); + MOCK_METHOD(void, UnregisterExternalKvByTier, (const std::string& node_id, TierType tier), + (override)); + MOCK_METHOD(void, UnregisterExternalKvByNode, (const std::string& node_id), (override)); + MOCK_METHOD(std::size_t, GarbageCollectHits, (std::chrono::system_clock::time_point cutoff), + (override)); + + // --- Block reads --- + MOCK_METHOD(std::vector, LookupBlock, (const std::string& key), (const, override)); + MOCK_METHOD(std::vector, LookupBlockForRouteGet, + (const std::string& key, (const std::unordered_set&)exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration), + (override)); + MOCK_METHOD(LocationBatch, BatchLookupBlockForRouteGet, + ((const std::vector&)keys, + (const std::unordered_set&)exclude_nodes, + std::chrono::system_clock::time_point now, + std::chrono::system_clock::duration lease_duration), + (override)); + MOCK_METHOD(std::vector, BatchExistsBlock, ((const std::vector&)keys), + (const, override)); + MOCK_METHOD((std::map>), EnumerateEvictionCandidates, + (const std::vector& buckets, EvictionOrder order, size_t max_per_bucket, + std::chrono::system_clock::time_point now), + (const, override)); + + // --- Client reads --- + MOCK_METHOD(std::optional, GetClient, (const std::string& node_id), + (const, override)); + MOCK_METHOD(bool, IsClientAlive, (const std::string& node_id), (const, override)); + MOCK_METHOD(std::optional, GetPeerAddress, (const std::string& node_id), + (const, override)); + MOCK_METHOD(std::vector, ListAliveClients, (), (const, override)); + MOCK_METHOD((std::unordered_map), GetAlivePeerView, (), + (const, override)); + MOCK_METHOD(std::size_t, AliveClientCount, (), (const, override)); + MOCK_METHOD(std::vector, GetClientTags, (const std::string& node_id), + (const, override)); + + // --- External-KV reads --- + MOCK_METHOD(std::vector, MatchExternalKv, + ((const std::vector&)hashes, bool count_as_hit, + std::chrono::system_clock::time_point now), + (override)); + MOCK_METHOD(std::vector, GetExternalKvHitCounts, + ((const std::vector&)hashes), (const, override)); + MOCK_METHOD(std::size_t, GetExternalKvCount, (const std::string& node_id), (const, override)); +}; + +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp b/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp new file mode 100644 index 000000000..29df4b893 --- /dev/null +++ b/tests/cpp/umbp/distributed/test_batch_resolve_codec.cpp @@ -0,0 +1,249 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Unit tests for the struct-of-arrays BatchResolveKeysResponse codec +// (batch_resolve_codec.h). This is the wire format shared by the peer service +// (encoder) and the pool client (decoder) after the resolve-keys payload +// optimization: per-key fields hoisted into parallel arrays, page locations +// flattened, buffer descriptors deduplicated once per batch, and page_size +// hoisted to a single field. The tests assert the encode->decode round-trip is +// lossless (correctness of the optimization), that the optimization actually +// shrinks the wire payload, that the omit_descs path is honored, and that a +// malformed response is rejected rather than over-read. +#include + +#include +#include +#include + +#include "umbp/distributed/peer/batch_resolve_codec.h" +#include "umbp/distributed/types.h" +#include "umbp_peer.pb.h" + +namespace mori::umbp { +namespace { + +BufferMemoryDescBytes MakeDesc(uint32_t buffer_index, const std::string& bytes) { + BufferMemoryDescBytes d; + d.buffer_index = buffer_index; + d.desc_bytes.assign(bytes.begin(), bytes.end()); + return d; +} + +// A representative batch: 4 keys mixing found / not-found, varying page counts +// (including a found-but-zero-page key), and several tiers. +std::vector SampleEntries() { + std::vector entries(4); + + entries[0].found = true; + entries[0].tier = TierType::HBM; + entries[0].size = 3 * 4096; + entries[0].pages = {{0, 10}, {1, 11}, {0, 12}}; + + entries[1].found = false; // miss: no tier / size / pages + + entries[2].found = true; + entries[2].tier = TierType::DRAM; + entries[2].size = 4096; + entries[2].pages = {{2, 7}}; + + entries[3].found = true; + entries[3].tier = TierType::SSD; + entries[3].size = 0; // found but zero pages (degenerate but legal) + entries[3].pages = {}; + + return entries; +} + +void ExpectKeyMatches(const ResolvedKeyEntry& src, const DecodedResolveKey& dec) { + EXPECT_EQ(src.found, dec.found); + if (!src.found) return; + EXPECT_EQ(src.tier, dec.tier); + EXPECT_EQ(src.size, dec.size); + ASSERT_EQ(src.pages.size(), dec.pages.size()); + for (size_t p = 0; p < src.pages.size(); ++p) { + EXPECT_EQ(src.pages[p].buffer_index, dec.pages[p].buffer_index); + EXPECT_EQ(src.pages[p].page_index, dec.pages[p].page_index); + } +} + +TEST(BatchResolveCodec, RoundTripPreservesEveryField) { + const auto entries = SampleEntries(); + const uint64_t page_size = 4096; + const std::vector descs = { + MakeDesc(0, "desc-zero"), MakeDesc(1, "desc-one"), MakeDesc(2, "desc-two")}; + + ::umbp::BatchResolveKeysResponse resp; + EncodeBatchResolveResponse(entries, page_size, descs, &resp); + + // Survives a real serialize/parse cycle, not just in-memory copy. + std::string wire; + ASSERT_TRUE(resp.SerializeToString(&wire)); + ::umbp::BatchResolveKeysResponse parsed; + ASSERT_TRUE(parsed.ParseFromString(wire)); + + EXPECT_EQ(BatchResolveKeyCount(parsed), static_cast(entries.size())); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(parsed); + EXPECT_EQ(decoded.page_size, page_size); + ASSERT_EQ(decoded.keys.size(), entries.size()); + for (size_t i = 0; i < entries.size(); ++i) ExpectKeyMatches(entries[i], decoded.keys[i]); + + ASSERT_EQ(decoded.descs.size(), descs.size()); + for (size_t i = 0; i < descs.size(); ++i) { + EXPECT_EQ(decoded.descs[i].buffer_index, descs[i].buffer_index); + EXPECT_EQ(decoded.descs[i].desc_bytes, descs[i].desc_bytes); + } +} + +TEST(BatchResolveCodec, FlattenedPagesSliceByKey) { + // Two found keys with different page counts must not bleed into each other. + std::vector entries(2); + entries[0].found = true; + entries[0].size = 2 * 4096; + entries[0].pages = {{5, 100}, {6, 101}}; + entries[1].found = true; + entries[1].size = 4096; + entries[1].pages = {{7, 200}}; + + ::umbp::BatchResolveKeysResponse resp; + EncodeBatchResolveResponse(entries, 4096, {}, &resp); + + // Flattened arrays hold exactly the concatenation, in order. + ASSERT_EQ(resp.buffer_index_size(), 3); + EXPECT_EQ(resp.buffer_index(0), 5u); + EXPECT_EQ(resp.buffer_index(1), 6u); + EXPECT_EQ(resp.buffer_index(2), 7u); + EXPECT_EQ(resp.page_count(0), 2u); + EXPECT_EQ(resp.page_count(1), 1u); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + ASSERT_EQ(decoded.keys.size(), 2u); + ASSERT_EQ(decoded.keys[0].pages.size(), 2u); + ASSERT_EQ(decoded.keys[1].pages.size(), 1u); + EXPECT_EQ(decoded.keys[1].pages[0].buffer_index, 7u); + EXPECT_EQ(decoded.keys[1].pages[0].page_index, 200u); +} + +TEST(BatchResolveCodec, OmitDescsRoundTrips) { + const auto entries = SampleEntries(); + + ::umbp::BatchResolveKeysResponse resp; + EncodeBatchResolveResponse(entries, 4096, /*descs=*/{}, &resp); + EXPECT_EQ(resp.descs_size(), 0); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + EXPECT_TRUE(decoded.descs.empty()); + ASSERT_EQ(decoded.keys.size(), entries.size()); + for (size_t i = 0; i < entries.size(); ++i) ExpectKeyMatches(entries[i], decoded.keys[i]); +} + +TEST(BatchResolveCodec, AllNotFound) { + std::vector entries(3); // all default => found=false + ::umbp::BatchResolveKeysResponse resp; + EncodeBatchResolveResponse(entries, 4096, {}, &resp); + + EXPECT_EQ(BatchResolveKeyCount(resp), 3); + EXPECT_EQ(resp.buffer_index_size(), 0); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + ASSERT_EQ(decoded.keys.size(), 3u); + for (const auto& k : decoded.keys) EXPECT_FALSE(k.found); +} + +TEST(BatchResolveCodec, EmptyBatch) { + ::umbp::BatchResolveKeysResponse resp; + EncodeBatchResolveResponse({}, 4096, {}, &resp); + EXPECT_EQ(BatchResolveKeyCount(resp), 0); + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + EXPECT_TRUE(decoded.keys.empty()); + EXPECT_EQ(decoded.page_size, 4096u); +} + +TEST(BatchResolveCodec, MalformedMismatchedArraysRejected) { + // found has 2 entries but the per-key arrays disagree -> decoder must refuse. + ::umbp::BatchResolveKeysResponse resp; + resp.set_page_size(4096); + resp.add_found(true); + resp.add_found(true); + resp.add_size(4096); // only one size for two found + resp.add_page_count(0); + resp.add_page_count(0); + // tier intentionally left empty (size 0 != 2) + + EXPECT_EQ(BatchResolveKeyCount(resp), 2); + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + EXPECT_TRUE(decoded.keys.empty()); // rejected, not partially decoded +} + +TEST(BatchResolveCodec, MalformedTruncatedPagesRejected) { + // page_count sums to 3 but only 1 flattened page is present. + ::umbp::BatchResolveKeysResponse resp; + resp.set_page_size(4096); + resp.add_found(true); + resp.add_tier(::umbp::TIER_HBM); + resp.add_size(3 * 4096); + resp.add_page_count(3); + resp.add_buffer_index(0); + resp.add_page_index(0); + + DecodedBatchResolve decoded = DecodeBatchResolveResponse(resp); + EXPECT_TRUE(decoded.keys.empty()); +} + +// The whole point of the optimization: with descriptors shared across many +// keys, the SoA layout that dedupes them once per batch is dramatically smaller +// than repeating a full descriptor set under every key. +TEST(BatchResolveCodec, OptimizedWireIsSmallerThanPerKeyDescs) { + const int kKeys = 1000; + const int kPagesPerKey = 8; + const int kBuffers = 4; + std::vector descs; + for (int b = 0; b < kBuffers; ++b) descs.push_back(MakeDesc(b, std::string(64, 'x'))); + + std::vector entries(kKeys); + for (int i = 0; i < kKeys; ++i) { + entries[i].found = true; + entries[i].tier = TierType::HBM; + entries[i].size = kPagesPerKey * 4096; + for (int p = 0; p < kPagesPerKey; ++p) { + entries[i].pages.push_back({static_cast(p % kBuffers), static_cast(p)}); + } + } + + ::umbp::BatchResolveKeysResponse opt; + EncodeBatchResolveResponse(entries, 4096, descs, &opt); + std::string opt_wire; + ASSERT_TRUE(opt.SerializeToString(&opt_wire)); + + // Naive baseline: the descriptor bytes repeated once per key (what the old + // per-entry ResolveKeyResponse layout carried). + const size_t per_key_desc_bytes = descs.size() * (64 + 4); + const size_t naive_desc_total = per_key_desc_bytes * kKeys; + const size_t opt_desc_total = per_key_desc_bytes; // once per batch + + EXPECT_LT(opt_wire.size(), naive_desc_total); + EXPECT_GT(naive_desc_total - opt_desc_total, 0u); +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp b/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp new file mode 100644 index 000000000..c5b47a67b --- /dev/null +++ b/tests/cpp/umbp/distributed/test_in_memory_master_metadata_store.cpp @@ -0,0 +1,897 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Phase 2 behavioral suite for InMemoryMasterMetadataStore (§6a). Written +// against IMasterMetadataStore& so the same cases validate the Redis backend +// later. Tests use injected system_clock times (no real-time sleeps) so they +// are deterministic in CI. +// +// State that the interface does not expose directly — lease_expiry and +// last_accessed_at on block entries — is observed through +// EnumerateEvictionCandidates: a leased entry is filtered out, and LRU +// ordering (EvictionOrder::kLeastRecentlyAccessed) reflects last_accessed_at. + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "umbp/distributed/master/in_memory_master_metadata_store.h" +#include "umbp/distributed/types.h" + +namespace mori::umbp { +namespace { + +using namespace std::chrono_literals; +using Clock = std::chrono::system_clock; + +// Fixed, NTP-plausible base instant so offsets read cleanly. +const Clock::time_point kT0 = Clock::time_point(std::chrono::hours(24 * 365 * 50)); + +std::map Caps(uint64_t total = 1000, uint64_t available = 1000) { + return {{TierType::HBM, TierCapacity{total, available}}}; +} + +ClientRegistration MakeReg(const std::string& node_id) { + ClientRegistration reg; + reg.node_id = node_id; + reg.node_address = "addr:" + node_id; + reg.peer_address = "peer:" + node_id; + reg.tier_capacities = Caps(); + reg.tags = {"role=test"}; + return reg; +} + +KvEvent Add(const std::string& key, TierType tier, uint64_t size) { + return KvEvent{KvEvent::Kind::ADD, key, tier, size}; +} +KvEvent Remove(const std::string& key, TierType tier) { + return KvEvent{KvEvent::Kind::REMOVE, key, tier, 0}; +} + +// Register `node` ALIVE at `now`. +void RegisterAlive(IMasterMetadataStore& store, const std::string& node, + Clock::time_point now = kT0) { + ASSERT_TRUE(store.RegisterClient(MakeReg(node), now, 30s)); +} + +// Apply a delta heartbeat carrying `events` at sequence `seq`. +HeartbeatResult Beat(IMasterMetadataStore& store, const std::string& node, uint64_t seq, + std::vector events, Clock::time_point now) { + return store.ApplyHeartbeat(node, seq, now, Caps(), events, /*is_full_sync=*/false); +} + +// --------------------------------------------------------------------------- +// RegisterClient +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, RegisterNewClient) { + InMemoryMasterMetadataStore store; + EXPECT_TRUE(store.RegisterClient(MakeReg("n1"), kT0, 30s)); + EXPECT_TRUE(store.IsClientAlive("n1")); + EXPECT_EQ(store.AliveClientCount(), 1u); + + auto rec = store.GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->status, ClientStatus::ALIVE); + EXPECT_EQ(rec->last_applied_seq, 0u); + EXPECT_EQ(rec->peer_address, "peer:n1"); + EXPECT_EQ(rec->last_heartbeat, kT0); + EXPECT_EQ(rec->registered_at, kT0); +} + +TEST(InMemoryStore, RejectReRegisterNonStaleAlive) { + InMemoryMasterMetadataStore store; + ASSERT_TRUE(store.RegisterClient(MakeReg("n1"), kT0, 30s)); + // Still well within stale_after window. + EXPECT_FALSE(store.RegisterClient(MakeReg("n1"), kT0 + 5s, 30s)); +} + +TEST(InMemoryStore, AcceptReRegisterStaleAlive) { + InMemoryMasterMetadataStore store; + ASSERT_TRUE(store.RegisterClient(MakeReg("n1"), kT0, 30s)); + // last_heartbeat is kT0; now - last_heartbeat > stale_after → re-register OK + // even though the reaper has not flipped the status yet (hazard #2). + EXPECT_TRUE(store.RegisterClient(MakeReg("n1"), kT0 + 31s, 30s)); + EXPECT_TRUE(store.IsClientAlive("n1")); +} + +TEST(InMemoryStore, AcceptReRegisterExpired) { + InMemoryMasterMetadataStore store; + ASSERT_TRUE(store.RegisterClient(MakeReg("n1"), kT0, 30s)); + ASSERT_EQ(store.ExpireStaleClients(kT0 + 1s).size(), 1u); + EXPECT_FALSE(store.IsClientAlive("n1")); + // Re-register an EXPIRED record at the same instant: accepted, back to ALIVE. + EXPECT_TRUE(store.RegisterClient(MakeReg("n1"), kT0 + 2s, 30s)); + EXPECT_TRUE(store.IsClientAlive("n1")); +} + +// --------------------------------------------------------------------------- +// UnregisterClient — cascade to block locations AND external KV +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, UnregisterClientCascades) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + + ASSERT_FALSE(store.LookupBlock("k1").empty()); + ASSERT_EQ(store.GetExternalKvCount("n1"), 2u); + + store.UnregisterClient("n1"); + + EXPECT_FALSE(store.GetClient("n1").has_value()); + EXPECT_TRUE(store.LookupBlock("k1").empty()); + EXPECT_EQ(store.GetExternalKvCount("n1"), 0u); + EXPECT_TRUE(store.MatchExternalKv({"h1", "h2"}, false, kT0).empty()); +} + +TEST(InMemoryStore, UnregisterUnknownIsNoOp) { + InMemoryMasterMetadataStore store; + store.UnregisterClient("ghost"); // must not crash + EXPECT_EQ(store.AliveClientCount(), 0u); +} + +// --------------------------------------------------------------------------- +// ApplyHeartbeat +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, HeartbeatUnknownNode) { + InMemoryMasterMetadataStore store; + auto r = Beat(store, "ghost", 1, {}, kT0); + EXPECT_EQ(r.status, HeartbeatResult::UNKNOWN); +} + +TEST(InMemoryStore, HeartbeatCasSequence) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + EXPECT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + EXPECT_EQ(Beat(store, "n1", 2, {Add("k2", TierType::HBM, 20)}, kT0).status, + HeartbeatResult::APPLIED); + // Out-of-order seq → SEQ_GAP, acked echoes last applied (2). + auto gap = Beat(store, "n1", 4, {Add("k3", TierType::HBM, 30)}, kT0); + EXPECT_EQ(gap.status, HeartbeatResult::SEQ_GAP); + EXPECT_EQ(gap.acked_seq, 2u); + // k3 must not have been applied. + EXPECT_TRUE(store.LookupBlock("k3").empty()); +} + +TEST(InMemoryStore, SeqGapKeepsLivenessNotCapsOrSeq) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {}, kT0).status, HeartbeatResult::APPLIED); + + // A gap heartbeat at a later time with different caps. + std::map new_caps = {{TierType::HBM, TierCapacity{9999, 9999}}}; + auto gap = store.ApplyHeartbeat("n1", 5, kT0 + 10s, new_caps, {}, /*is_full_sync=*/false); + ASSERT_EQ(gap.status, HeartbeatResult::SEQ_GAP); + + auto rec = store.GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->status, ClientStatus::ALIVE); // kept alive + EXPECT_EQ(rec->last_heartbeat, kT0 + 10s); // last_heartbeat bumped + EXPECT_EQ(rec->last_applied_seq, 1u); // seq NOT advanced + EXPECT_EQ(rec->tier_capacities.at(TierType::HBM).total_bytes, 1000u); // caps NOT replaced +} + +TEST(InMemoryStore, HeartbeatDeltaAddRemove) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_EQ(store.LookupBlock("k1").size(), 1u); + + ASSERT_EQ(Beat(store, "n1", 2, {Remove("k1", TierType::HBM)}, kT0).status, + HeartbeatResult::APPLIED); + EXPECT_TRUE(store.LookupBlock("k1").empty()); +} + +TEST(InMemoryStore, HeartbeatFullSyncReplaces) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10), Add("k2", TierType::HBM, 20)}, kT0) + .status, + HeartbeatResult::APPLIED); + + // full_sync wipes prior locations and installs only the ADDs carried here. + auto r = store.ApplyHeartbeat("n1", 7, kT0, Caps(), {Add("k3", TierType::HBM, 30)}, + /*is_full_sync=*/true); + EXPECT_EQ(r.status, HeartbeatResult::APPLIED); + EXPECT_EQ(r.acked_seq, 7u); + + EXPECT_TRUE(store.LookupBlock("k1").empty()); + EXPECT_TRUE(store.LookupBlock("k2").empty()); + EXPECT_EQ(store.LookupBlock("k3").size(), 1u); + + auto rec = store.GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->last_applied_seq, 7u); // full_sync re-baselines the seq +} + +// --------------------------------------------------------------------------- +// ExpireStaleClients — flip to EXPIRED, keep row, cascade, idempotent +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, ExpireStaleFlipsKeepsRowAndCascades) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1", kT0); + RegisterAlive(store, "n2", kT0 + 20s); // fresher + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + + // Cutoff after n1's heartbeat but before n2's. + auto dead = store.ExpireStaleClients(kT0 + 10s); + ASSERT_EQ(dead.size(), 1u); + EXPECT_EQ(dead[0], "n1"); + + // Row KEPT but EXPIRED (hazard #3). + auto rec = store.GetClient("n1"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->status, ClientStatus::EXPIRED); + EXPECT_FALSE(store.IsClientAlive("n1")); + + // Cascade dropped its blocks and external KV. + EXPECT_TRUE(store.LookupBlock("k1").empty()); + EXPECT_EQ(store.GetExternalKvCount("n1"), 0u); + + // n2 untouched. + EXPECT_TRUE(store.IsClientAlive("n2")); +} + +TEST(InMemoryStore, ExpireStaleIsIdempotent) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1", kT0); + ASSERT_EQ(store.ExpireStaleClients(kT0 + 10s).size(), 1u); + // Re-tick: already EXPIRED, nothing new to report. + EXPECT_TRUE(store.ExpireStaleClients(kT0 + 10s).empty()); +} + +TEST(InMemoryStore, ExpiredRowExcludedFromAliveAccounting) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1", kT0); + RegisterAlive(store, "n2", kT0); + ASSERT_EQ(store.AliveClientCount(), 2u); + ASSERT_EQ(store.ExpireStaleClients(kT0 + 10s).size(), 2u); + + EXPECT_EQ(store.AliveClientCount(), 0u); // not 2, even though rows remain + EXPECT_TRUE(store.ListAliveClients().empty()); + EXPECT_TRUE(store.GetClient("n1").has_value()); // row still present +} + +// --------------------------------------------------------------------------- +// Block reads — lease/access observed via EnumerateEvictionCandidates +// --------------------------------------------------------------------------- + +// Helper: single (node, tier) bucket to enumerate candidates from. The store's +// candidate enumeration is policy-neutral: it takes the buckets to scan, an +// ordering hint, and a per-bucket count cap (0 = no cap). The byte-budget / +// victim decision lives in MasterEvictStrategy, not the store. +std::vector Buckets(const std::string& node, TierType tier) { + return {NodeTierKey{node, tier}}; +} + +TEST(InMemoryStore, LookupBlockHasNoLeaseOrAccessSideEffects) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + + // Plain read twice. + EXPECT_EQ(store.LookupBlock("k1").size(), 1u); + EXPECT_EQ(store.LookupBlock("k1").size(), 1u); + + // Not leased → still an eviction candidate at kT0. + auto cands = store.EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0); + ASSERT_EQ(cands.size(), 1u); + EXPECT_EQ(cands.begin()->second.size(), 1u); +} + +TEST(InMemoryStore, LookupBlockForRouteGetGrantsLeaseAndAccess) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + + auto locs = store.LookupBlockForRouteGet("k1", {}, kT0, 60s); + ASSERT_EQ(locs.size(), 1u); + + // Leased until kT0+60s → filtered out of eviction at kT0+10s. + EXPECT_TRUE(store + .EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0 + 10s) + .empty()); + // After lease expiry it is a candidate again. + EXPECT_FALSE(store + .EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0 + 61s) + .empty()); +} + +TEST(InMemoryStore, RouteGetExcludeNodesNoLeaseWhenFullyExcluded) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + + std::unordered_set exclude = {"n1"}; + auto locs = store.LookupBlockForRouteGet("k1", exclude, kT0, 60s); + EXPECT_TRUE(locs.empty()); // every location excluded + + // No lease granted (hazard #4) → still an eviction candidate immediately. + EXPECT_FALSE(store + .EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0) + .empty()); +} + +TEST(InMemoryStore, BatchLookupForRouteGetParallelToKeys) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10), Add("k3", TierType::HBM, 30)}, kT0) + .status, + HeartbeatResult::APPLIED); + + auto out = store.BatchLookupBlockForRouteGet({"k1", "missing", "k3"}, {}, kT0, 60s); + ASSERT_EQ(out.size(), 3u); + EXPECT_EQ(out[0].size(), 1u); + EXPECT_TRUE(out[1].empty()); // missing key + EXPECT_EQ(out[2].size(), 1u); +} + +TEST(InMemoryStore, BatchExistsBlockNoSideEffects) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + + auto exists = store.BatchExistsBlock({"k1", "missing"}); + ASSERT_EQ(exists.size(), 2u); + EXPECT_TRUE(exists[0]); + EXPECT_FALSE(exists[1]); + + // No lease granted by an existence check. + EXPECT_FALSE(store + .EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0) + .empty()); +} + +// --------------------------------------------------------------------------- +// EnumerateEvictionCandidates +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, EvictionLruOrderAndCap) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + // Three keys, each 100 bytes, accessed at increasing times so LRU order is + // k_old < k_mid < k_new. + ASSERT_EQ(Beat(store, "n1", 1, {Add("k_old", TierType::HBM, 100)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_EQ(Beat(store, "n1", 2, {Add("k_mid", TierType::HBM, 100)}, kT0 + 1s).status, + HeartbeatResult::APPLIED); + ASSERT_EQ(Beat(store, "n1", 3, {Add("k_new", TierType::HBM, 100)}, kT0 + 2s).status, + HeartbeatResult::APPLIED); + + // max_per_bucket=2 with LRU order → the two oldest, oldest first. (The store + // caps by count; trimming to a byte budget is MasterEvictStrategy's job.) + auto cands = store.EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, + /*max_per_bucket=*/2, kT0 + 10s); + ASSERT_EQ(cands.size(), 1u); + auto& bucket = cands.at(NodeTierKey{"n1", TierType::HBM}); + ASSERT_EQ(bucket.size(), 2u); + EXPECT_EQ(bucket[0].key, "k_old"); // oldest first + EXPECT_EQ(bucket[1].key, "k_mid"); +} + +TEST(InMemoryStore, EvictionSkipsLeased) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 100)}, kT0).status, + HeartbeatResult::APPLIED); + // Lease k1 well past the enumeration time. + store.LookupBlockForRouteGet("k1", {}, kT0, 1h); + EXPECT_TRUE(store + .EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, kT0 + 1s) + .empty()); +} + +TEST(InMemoryStore, EvictionTieTimestampsAllSurvive) { + // §2d correctness claim: many candidates sharing one identical last_accessed_at + // (the common case, since a batch RouteGet stamps one `now` across all keys) + // must all be enumerable — none dropped by tie collisions. + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + std::vector adds; + for (int i = 0; i < 50; ++i) { + adds.push_back(Add("k" + std::to_string(i), TierType::HBM, 10)); + } + // All keys created (and thus last_accessed) at the identical instant kT0. + ASSERT_EQ(Beat(store, "n1", 1, adds, kT0).status, HeartbeatResult::APPLIED); + + // No cap → take everything; all 50 tied-timestamp candidates must appear. + auto cands = store.EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, + /*max_per_bucket=*/0, kT0 + 10s); + ASSERT_EQ(cands.size(), 1u); + EXPECT_EQ(cands.at(NodeTierKey{"n1", TierType::HBM}).size(), 50u); +} + +TEST(InMemoryStore, EvictionOnlyRequestedBuckets) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("kh", TierType::HBM, 10), Add("kd", TierType::DRAM, 10)}, kT0) + .status, + HeartbeatResult::APPLIED); + // Only ask about the HBM bucket. + auto cands = store.EnumerateEvictionCandidates( + Buckets("n1", TierType::HBM), EvictionOrder::kLeastRecentlyAccessed, 0, kT0 + 1s); + ASSERT_EQ(cands.size(), 1u); + EXPECT_EQ(cands.begin()->first.tier, TierType::HBM); +} + +// --------------------------------------------------------------------------- +// External KV +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, RegisterExternalKvAliveGate) { + InMemoryMasterMetadataStore store; + // Dead/unknown node → rejected, nothing written. + EXPECT_FALSE(store.RegisterExternalKvIfAlive("ghost", {"h1"}, TierType::HBM)); + EXPECT_TRUE(store.MatchExternalKv({"h1"}, false, kT0).empty()); + + RegisterAlive(store, "n1"); + EXPECT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + EXPECT_EQ(store.MatchExternalKv({"h1"}, false, kT0).size(), 1u); +} + +TEST(InMemoryStore, UnregisterExternalKvAndByTier) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + + // Remove only the HBM tier; DRAM remains. + store.UnregisterExternalKv("n1", {"h1"}, TierType::HBM); + auto m = store.MatchExternalKv({"h1"}, false, kT0); + ASSERT_EQ(m.size(), 1u); + EXPECT_EQ(m[0].hashes_by_tier.count(TierType::HBM), 0u); + EXPECT_EQ(m[0].hashes_by_tier.count(TierType::DRAM), 1u); + + // Whole-tier wipe of DRAM → entry gone. + store.UnregisterExternalKvByTier("n1", TierType::DRAM); + EXPECT_TRUE(store.MatchExternalKv({"h1"}, false, kT0).empty()); +} + +TEST(InMemoryStore, MatchCountsHitsWhenRequested) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + + // count_as_hit=false: pure read, hit map untouched. + store.MatchExternalKv({"h1", "h2"}, /*count_as_hit=*/false, kT0); + EXPECT_TRUE(store.GetExternalKvHitCounts({"h1", "h2"}).empty()); + + // count_as_hit=true: increments accumulate across calls. + store.MatchExternalKv({"h1", "h2"}, /*count_as_hit=*/true, kT0); + store.MatchExternalKv({"h1"}, /*count_as_hit=*/true, kT0 + 1s); + + auto counts = store.GetExternalKvHitCounts({"h1", "h2"}); + std::map by_hash; + for (const auto& e : counts) by_hash[e.hash] = e.hit_count_total; + EXPECT_EQ(by_hash["h1"], 2u); + EXPECT_EQ(by_hash["h2"], 1u); +} + +TEST(InMemoryStore, MatchedHashCountAcrossTiers) { + // Preserves the NodeMatch::MatchedHashCount coverage from + // test_external_kv_block_index.cpp:57 — one hash mirrored across two tiers + // counts once. + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + + auto m = store.MatchExternalKv({"h1"}, false, kT0); + ASSERT_EQ(m.size(), 1u); + EXPECT_EQ(m[0].hashes_by_tier.size(), 2u); // appears in two tier buckets + EXPECT_EQ(m[0].MatchedHashCount(), 1u); // but is one unique hash +} + +TEST(InMemoryStore, GetExternalKvHitCountsDedupesAndSkipsMissing) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::HBM)); + store.MatchExternalKv({"h1"}, true, kT0); + + auto counts = store.GetExternalKvHitCounts({"missing", "h1", "h1"}); + ASSERT_EQ(counts.size(), 1u); + EXPECT_EQ(counts[0].hash, "h1"); + EXPECT_EQ(counts[0].hit_count_total, 1u); +} + +TEST(InMemoryStore, GarbageCollectHitsByLastSeen) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"old", "fresh"}, TierType::HBM)); + store.MatchExternalKv({"old"}, true, kT0); + store.MatchExternalKv({"fresh"}, true, kT0 + 100s); + + // Drop entries last seen before kT0+50s → only "old" goes. + EXPECT_EQ(store.GarbageCollectHits(kT0 + 50s), 1u); + + auto counts = store.GetExternalKvHitCounts({"old", "fresh"}); + ASSERT_EQ(counts.size(), 1u); + EXPECT_EQ(counts[0].hash, "fresh"); +} + +TEST(InMemoryStore, UnregisterExternalKvByNodeWipesAllTiersOnly) { + // Whole-node external-KV wipe (backs RevokeAllExternalKvBlocksForNode). Unlike + // UnregisterClient, it must NOT touch the client record or block locations. + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + ASSERT_EQ(Beat(store, "n1", 1, {Add("k1", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1", "h2"}, TierType::HBM)); + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1"}, TierType::DRAM)); + ASSERT_EQ(store.GetExternalKvCount("n1"), 2u); + + store.UnregisterExternalKvByNode("n1"); + + // External KV gone across every tier. + EXPECT_EQ(store.GetExternalKvCount("n1"), 0u); + EXPECT_TRUE(store.MatchExternalKv({"h1", "h2"}, false, kT0).empty()); + + // Client record and block locations untouched (distinguishes from UnregisterClient). + EXPECT_TRUE(store.IsClientAlive("n1")); + EXPECT_EQ(store.LookupBlock("k1").size(), 1u); +} + +TEST(InMemoryStore, UnregisterExternalKvByNodeUnknownIsNoOp) { + InMemoryMasterMetadataStore store; + store.UnregisterExternalKvByNode("ghost"); // must not crash + EXPECT_EQ(store.GetExternalKvCount("ghost"), 0u); +} + +// --------------------------------------------------------------------------- +// Client reads — GetPeerAddress, GetClientTags, ListAliveClients content +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, GetPeerAddressAliveExpiredAndUnknown) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + + // ALIVE → peer surfaced (MakeReg sets peer:). + auto alive = store.GetPeerAddress("n1"); + ASSERT_TRUE(alive.has_value()); + EXPECT_EQ(*alive, "peer:n1"); + + // EXPIRED rows still surface their peer_address (contract: the row is kept). + ASSERT_EQ(store.ExpireStaleClients(kT0 + 10s).size(), 1u); + auto expired = store.GetPeerAddress("n1"); + ASSERT_TRUE(expired.has_value()); + EXPECT_EQ(*expired, "peer:n1"); + + // Unknown node → nullopt. + EXPECT_FALSE(store.GetPeerAddress("ghost").has_value()); +} + +TEST(InMemoryStore, GetClientTagsReturnsRegisteredTagsAndEmptyForUnknown) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); // MakeReg sets tags = {"role=test"} + + auto tags = store.GetClientTags("n1"); + ASSERT_EQ(tags.size(), 1u); + EXPECT_EQ(tags[0], "role=test"); + + EXPECT_TRUE(store.GetClientTags("ghost").empty()); +} + +TEST(InMemoryStore, ListAliveClientsReturnsAliveRecordsExcludingExpired) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1", kT0); + RegisterAlive(store, "n2", kT0 + 20s); // fresher, survives the cutoff below + + // Expire only n1. + ASSERT_EQ(store.ExpireStaleClients(kT0 + 10s).size(), 1u); + + auto alive = store.ListAliveClients(); + ASSERT_EQ(alive.size(), 1u); // n1 excluded even though its row still exists + EXPECT_EQ(alive[0].node_id, "n2"); + EXPECT_EQ(alive[0].status, ClientStatus::ALIVE); + EXPECT_EQ(alive[0].peer_address, "peer:n2"); +} + +// The lightweight peer view maps node->peer (no capacity) and reflects +// membership changes. Ported from the old ClientRegistry peer-view tests. +TEST(InMemoryStore, GetAlivePeerViewTracksMembership) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + + auto v1 = store.GetAlivePeerView(); + EXPECT_EQ(v1.size(), 1u); + ASSERT_EQ(v1.count("n1"), 1u); + EXPECT_EQ(v1.at("n1"), "peer:n1"); // MakeReg sets peer: + + // Membership change → reflected in a freshly built view. + RegisterAlive(store, "n2"); + auto v2 = store.GetAlivePeerView(); + EXPECT_EQ(v2.size(), 2u); + EXPECT_EQ(v2.at("n2"), "peer:n2"); +} + +// EXPIRED rows are excluded from the peer view (unlike GetPeerAddress, which +// still surfaces a single expired row's peer). +TEST(InMemoryStore, GetAlivePeerViewExcludesExpired) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1", kT0); + RegisterAlive(store, "n2", kT0 + 20s); // fresher, survives the cutoff below + + ASSERT_EQ(store.ExpireStaleClients(kT0 + 10s).size(), 1u); // expires n1 only + + auto view = store.GetAlivePeerView(); + EXPECT_EQ(view.size(), 1u); + EXPECT_EQ(view.count("n1"), 0u); + ASSERT_EQ(view.count("n2"), 1u); + EXPECT_EQ(view.at("n2"), "peer:n2"); +} + +// The peer view carries no capacity, so a capacity-only heartbeat leaves its +// contents (node → peer) unchanged. +TEST(InMemoryStore, GetAlivePeerViewIgnoresCapacity) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + + auto p1 = store.GetAlivePeerView(); + ASSERT_EQ(Beat(store, "n1", /*seq=*/1, /*events=*/{}, kT0 + 1s).status, HeartbeatResult::APPLIED); + auto p2 = store.GetAlivePeerView(); + EXPECT_EQ(p1, p2); +} + +// --------------------------------------------------------------------------- +// Concurrency +// --------------------------------------------------------------------------- + +TEST(InMemoryStore, ConcurrentHeartbeatCasExactlyOneApplied) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + + std::atomic applied{0}; + std::atomic gap{0}; + std::atomic start{false}; + std::vector threads; + for (int t = 0; t < 2; ++t) { + threads.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + // Both race to apply seq=1 (last_applied starts at 0). + auto r = store.ApplyHeartbeat("n1", 1, kT0, Caps(), {}, /*is_full_sync=*/false); + if (r.status == HeartbeatResult::APPLIED) { + applied.fetch_add(1); + } else if (r.status == HeartbeatResult::SEQ_GAP) { + gap.fetch_add(1); + } + }); + } + start.store(true, std::memory_order_release); + for (auto& th : threads) th.join(); + + EXPECT_EQ(applied.load(), 1); + EXPECT_EQ(gap.load(), 1); + EXPECT_EQ(store.GetClient("n1")->last_applied_seq, 1u); +} + +// ThreadSanitizer safety net for collapsing four lock domains into one: a mixed +// read/write workload across the shared/unique split must be race-free. +TEST(InMemoryStore, MixedWorkloadIsRaceFree) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "n1"); + for (int i = 0; i < 100; ++i) { + store.ApplyHeartbeat("n1", i + 1, kT0, Caps(), + {Add("k" + std::to_string(i), TierType::HBM, 10)}, + /*is_full_sync=*/false); + } + ASSERT_TRUE(store.RegisterExternalKvIfAlive("n1", {"h1", "h2", "h3"}, TierType::HBM)); + + std::atomic start{false}; + std::vector threads; + + // RouteGet readers (shared-lock path with atomic lease/access mutation). + for (int r = 0; r < 4; ++r) { + threads.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + for (int i = 0; i < 500; ++i) { + store.BatchLookupBlockForRouteGet({"k1", "k50", "k99"}, {}, kT0 + std::chrono::seconds(i), + 30s); + store.BatchExistsBlock({"k1", "k2"}); + } + }); + } + // Hit writers (the formerly-shared path that becomes exclusive). + threads.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + for (int i = 0; i < 500; ++i) { + store.MatchExternalKv({"h1", "h2", "h3"}, /*count_as_hit=*/true, + kT0 + std::chrono::seconds(i)); + } + }); + // Eviction-enumeration reader. + threads.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + for (int i = 0; i < 500; ++i) { + store.EnumerateEvictionCandidates(Buckets("n1", TierType::HBM), + EvictionOrder::kLeastRecentlyAccessed, 0, + kT0 + std::chrono::seconds(i)); + } + }); + + start.store(true, std::memory_order_release); + for (auto& th : threads) th.join(); + + // After the storm, hit counts reflect exactly the 500 hit-writer iterations. + auto counts = store.GetExternalKvHitCounts({"h1"}); + ASSERT_EQ(counts.size(), 1u); + EXPECT_EQ(counts[0].hit_count_total, 500u); +} + +// --------------------------------------------------------------------------- +// Cross-domain atomicity conformance (approach A). These make the intermediate- +// state declarations in master_metadata_store.h executable: with the block +// index split into key-hashed shards and cross-domain writes running +// meta→block WITHOUT nesting the locks, a concurrent RouteGet reader must still +// (a) never observe a torn per-key location, and (b) never resolve a peer for a +// node that isn't in the alive-peer view (residual block locations for an +// erased node are simply skipped by the router-style join — router.cpp). +// --------------------------------------------------------------------------- + +// A writer repeatedly full-syncs a node's whole key set between two well-formed +// shapes (size 100 vs size 200). full_sync clears+replays per shard; within a +// shard a single key's replacement is atomic, so a reader must always see an +// old-or-new value, never a torn one. Also a ThreadSanitizer net for the +// meta→block split under concurrent full_sync. +TEST(InMemoryStore, ConcurrentFullSyncRouteGetNeverTears) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "w"); + + const int kNumKeys = 64; // spread across shards (default 32) + std::vector keys; + keys.reserve(kNumKeys); + std::vector adds100, adds200; + for (int i = 0; i < kNumKeys; ++i) { + keys.push_back("k" + std::to_string(i)); + adds100.push_back(Add(keys.back(), TierType::HBM, 100)); + adds200.push_back(Add(keys.back(), TierType::HBM, 200)); + } + ASSERT_EQ(store.ApplyHeartbeat("w", 1, kT0, Caps(), adds100, /*is_full_sync=*/false).status, + HeartbeatResult::APPLIED); + + std::atomic start{false}; + std::atomic stop{false}; + std::atomic torn{false}; + + std::thread writer([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + uint64_t seq = 100; // full_sync re-baselines seq, so any increasing value is fine + for (int it = 0; it < 2000; ++it) { + const auto& adds = (it % 2 == 0) ? adds200 : adds100; + store.ApplyHeartbeat("w", seq++, kT0, Caps(), adds, /*is_full_sync=*/true); + } + stop.store(true, std::memory_order_release); + }); + + std::vector readers; + for (int r = 0; r < 4; ++r) { + readers.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + while (!stop.load(std::memory_order_acquire)) { + auto locs = store.BatchLookupBlockForRouteGet(keys, {}, kT0, 30s); + for (const auto& per_key : locs) { + for (const auto& loc : per_key) { + const bool well_formed = loc.node_id == "w" && loc.tier == TierType::HBM && + (loc.size == 100 || loc.size == 200); + if (!well_formed) torn.store(true, std::memory_order_release); + } + } + } + }); + } + start.store(true, std::memory_order_release); + writer.join(); + for (auto& t : readers) t.join(); + EXPECT_FALSE(torn.load()); // every visible location was a valid old-or-new value +} + +// A writer churns node "b" (Unregister → Register → re-seed key K) while readers +// do a router-style RouteGet: snapshot GetAlivePeerView, then join it with the +// block locations. The invariant: a location whose node resolves in the alive +// view always yields a non-empty peer for a node that WAS alive in that same +// snapshot — a residual block location for an erased node is absent from the +// view and thus never routed to. Exercises the meta→block non-atomic window. +TEST(InMemoryStore, ConcurrentUnregisterRouteGetNeverResolvesDeadPeer) { + InMemoryMasterMetadataStore store; + RegisterAlive(store, "a"); + RegisterAlive(store, "b"); + ASSERT_EQ(Beat(store, "a", 1, {Add("K", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + ASSERT_EQ(Beat(store, "b", 1, {Add("K", TierType::HBM, 10)}, kT0).status, + HeartbeatResult::APPLIED); + + std::atomic start{false}; + std::atomic bad_peer{false}; + + std::thread writer([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + for (int it = 0; it < 1000; ++it) { + store.UnregisterClient("b"); // meta erase + block wipe + store.RegisterClient(MakeReg("b"), kT0, 30s); // back to ALIVE + store.ApplyHeartbeat("b", 1, kT0, Caps(), {Add("K", TierType::HBM, 10)}, + /*is_full_sync=*/true); // re-seed K@b + } + }); + + std::vector readers; + for (int r = 0; r < 4; ++r) { + readers.emplace_back([&] { + while (!start.load(std::memory_order_acquire)) std::this_thread::yield(); + for (int it = 0; it < 2000; ++it) { + // Router-style resolution (router.cpp): peer comes from the alive view, + // NOT from the block location itself. + auto view = store.GetAlivePeerView(); + auto locs = store.BatchLookupBlockForRouteGet({"K"}, {}, kT0, 30s); + for (const auto& loc : locs[0]) { + auto pit = view.find(loc.node_id); + if (pit != view.end() && pit->second.empty()) { + bad_peer.store(true, std::memory_order_release); + } + } + } + }); + } + start.store(true, std::memory_order_release); + writer.join(); + for (auto& t : readers) t.join(); + EXPECT_FALSE(bad_peer.load()); + + // Deterministic post-condition: with b unregistered for good, any location we + // could route to (i.e. that resolves in the alive view) is a live node, never + // a stale b; b's residual location (if still present) is absent from the view. + store.UnregisterClient("b"); + auto view = store.GetAlivePeerView(); + auto locs = store.BatchLookupBlockForRouteGet({"K"}, {}, kT0, 30s); + EXPECT_EQ(view.count("b"), 0u); + for (const auto& loc : locs[0]) { + if (view.count(loc.node_id)) EXPECT_EQ(loc.node_id, "a"); + } +} + +} // namespace +} // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp b/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp index 97c29e12a..9fed02443 100644 --- a/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_master_evict_strategy.cpp @@ -32,13 +32,13 @@ #include #include "umbp/distributed/master/evict_strategy.h" -#include "umbp/distributed/master/global_block_index.h" #include "umbp/distributed/types.h" namespace mori::umbp { namespace { -using Clock = std::chrono::steady_clock; +// EvictionCandidate::last_accessed_at is a system_clock time_point. +using Clock = std::chrono::system_clock; EvictionCandidate MakeCandidate(const std::string& key, const std::string& node, TierType tier, uint64_t size, Clock::time_point accessed) { diff --git a/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp b/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp new file mode 100644 index 000000000..9fcd952cd --- /dev/null +++ b/tests/cpp/umbp/distributed/test_master_metadata_store_interface.cpp @@ -0,0 +1,134 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// Phase 1 compile/instantiation gate for IMasterMetadataStore. +// +// The interface is abstract with no implementation yet, so there is no runtime +// behavior to exercise. The bar for Phase 1 is that the contract is well-formed +// and instantiable: +// 1. MockMasterMetadataStore overrides every pure-virtual (a missing or +// ill-typed override makes the mock abstract → fails to instantiate). +// 2. A MockMasterMetadataStore is usable through an IMasterMetadataStore&, +// proving the override set is complete. +// Behavioral assertions arrive with InMemoryMasterMetadataStore in Phase 2. + +#include + +#include + +#include "mock_master_metadata_store.h" +#include "umbp/distributed/master/master_metadata_store.h" + +namespace mori::umbp { +namespace { + +// Instantiation gate: if the interface had an orphaned or ill-typed pure +// virtual, MockMasterMetadataStore would stay abstract and this would not +// compile. +TEST(MasterMetadataStoreInterface, MockIsInstantiableThroughInterface) { + MockMasterMetadataStore mock; + IMasterMetadataStore& store = mock; + (void)store; + SUCCEED(); +} + +// Signature-completeness spot check: name every interface method once through +// the base-class pointer with a default ON_CALL, mirroring the §1b delta table +// plus the two added hit-count methods (GetExternalKvHitCounts, +// GarbageCollectHits) and the `now` parameter on MatchExternalKv. This guards +// against silently dropping the live GetExternalKvHitCounts RPC path. +TEST(MasterMetadataStoreInterface, EveryMethodIsCallableThroughInterface) { + using ::testing::_; + using ::testing::NiceMock; + using ::testing::Return; + using namespace std::chrono_literals; + + // NiceMock: these are default-action calls, not behavior under test, so the + // "uninteresting call" warnings would just be noise. + NiceMock mock; + const auto now = std::chrono::system_clock::now(); + + ON_CALL(mock, RegisterClient(_, _, _)).WillByDefault(Return(true)); + ON_CALL(mock, ApplyHeartbeat(_, _, _, _, _, _)) + .WillByDefault(Return(HeartbeatResult{HeartbeatResult::APPLIED, 0})); + ON_CALL(mock, ExpireStaleClients(_)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, RegisterExternalKvIfAlive(_, _, _)).WillByDefault(Return(true)); + ON_CALL(mock, GarbageCollectHits(_)).WillByDefault(Return(0)); + ON_CALL(mock, LookupBlock(_)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, LookupBlockForRouteGet(_, _, _, _)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, BatchLookupBlockForRouteGet(_, _, _, _)) + .WillByDefault(Return(std::vector>{})); + ON_CALL(mock, BatchExistsBlock(_)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, EnumerateEvictionCandidates(_, _, _, _)) + .WillByDefault(Return(std::map>{})); + ON_CALL(mock, GetClient(_)).WillByDefault(Return(std::nullopt)); + ON_CALL(mock, IsClientAlive(_)).WillByDefault(Return(false)); + ON_CALL(mock, GetPeerAddress(_)).WillByDefault(Return(std::nullopt)); + ON_CALL(mock, ListAliveClients()).WillByDefault(Return(std::vector{})); + ON_CALL(mock, AliveClientCount()).WillByDefault(Return(0)); + ON_CALL(mock, GetClientTags(_)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, MatchExternalKv(_, _, _)).WillByDefault(Return(std::vector{})); + ON_CALL(mock, GetExternalKvHitCounts(_)) + .WillByDefault(Return(std::vector{})); + ON_CALL(mock, GetExternalKvCount(_)).WillByDefault(Return(0)); + + IMasterMetadataStore& store = mock; + + // Cross-store writes. + ClientRegistration reg; + reg.node_id = "node-a"; + EXPECT_TRUE(store.RegisterClient(reg, now, 30s)); + store.UnregisterClient("node-a"); + EXPECT_EQ(store.ApplyHeartbeat("node-a", 1, now, {}, {}, false).status, HeartbeatResult::APPLIED); + EXPECT_TRUE(store.ExpireStaleClients(now).empty()); + + // External-KV writes. + EXPECT_TRUE(store.RegisterExternalKvIfAlive("node-a", {"h0"}, TierType::HBM)); + store.UnregisterExternalKv("node-a", {"h0"}, TierType::HBM); + store.UnregisterExternalKvByTier("node-a", TierType::HBM); + store.UnregisterExternalKvByNode("node-a"); + EXPECT_EQ(store.GarbageCollectHits(now), 0u); + + // Block reads. + EXPECT_TRUE(store.LookupBlock("k0").empty()); + EXPECT_TRUE(store.LookupBlockForRouteGet("k0", {}, now, 5s).empty()); + EXPECT_TRUE(store.BatchLookupBlockForRouteGet({"k0"}, {}, now, 5s).empty()); + EXPECT_TRUE(store.BatchExistsBlock({"k0"}).empty()); + EXPECT_TRUE( + store.EnumerateEvictionCandidates({}, EvictionOrder::kLeastRecentlyAccessed, 0, now).empty()); + + // Client reads. + EXPECT_FALSE(store.GetClient("node-a").has_value()); + EXPECT_FALSE(store.IsClientAlive("node-a")); + EXPECT_FALSE(store.GetPeerAddress("node-a").has_value()); + EXPECT_TRUE(store.ListAliveClients().empty()); + EXPECT_EQ(store.AliveClientCount(), 0u); + EXPECT_TRUE(store.GetClientTags("node-a").empty()); + + // External-KV reads, incl. the two added hit-count methods + `now` param. + EXPECT_TRUE(store.MatchExternalKv({"h0"}, /*count_as_hit=*/true, now).empty()); + EXPECT_TRUE(store.GetExternalKvHitCounts({"h0"}).empty()); + EXPECT_EQ(store.GetExternalKvCount("node-a"), 0u); +} + +} // namespace +} // namespace mori::umbp diff --git a/src/umbp/tests/test_peer_dram_allocator.cpp b/tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp similarity index 97% rename from src/umbp/tests/test_peer_dram_allocator.cpp rename to tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp index dda9d4c82..da2ac1a66 100644 --- a/src/umbp/tests/test_peer_dram_allocator.cpp +++ b/tests/cpp/umbp/distributed/test_peer_dram_allocator.cpp @@ -517,30 +517,6 @@ TEST(PeerDramAllocator, BatchResolveMixedHitsAndMisses) { EXPECT_FALSE(results[3].found); } -TEST(PeerDramAllocator, BatchResolveIncludeDescsFalseSkipsDescs) { - auto a = MakeAllocator(); - auto p_hit = AllocateOk(*a, "hit", kPageSize * 5, TierType::DRAM); - ASSERT_TRUE(p_hit.has_value()); - uint64_t committed_bytes = 0; - ASSERT_TRUE(a->Commit(p_hit->slot_id, "hit", committed_bytes)); - a->DrainPendingEvents(); - - auto ref_hit = a->Resolve("hit"); - ASSERT_TRUE(ref_hit.found); - - auto results = a->BatchResolve({"hit", "ghost"}, /*include_descs=*/false); - ASSERT_EQ(results.size(), 2u); - - EXPECT_TRUE(results[0].found); - EXPECT_EQ(results[0].tier, ref_hit.tier); - EXPECT_EQ(results[0].pages, ref_hit.pages); - EXPECT_EQ(results[0].size, ref_hit.size); - EXPECT_TRUE(results[0].descs.empty()); - - EXPECT_FALSE(results[1].found); - EXPECT_TRUE(results[1].descs.empty()); -} - TEST(PeerDramAllocator, BatchResolveExtendsLeaseForHitsOnly) { auto a = std::make_unique(kPageSize, MakeDramCfg(), EmptyCfg(), /*pending_ttl=*/std::chrono::milliseconds{5000}, diff --git a/src/umbp/tests/test_peer_ssd_eviction.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp similarity index 100% rename from src/umbp/tests/test_peer_ssd_eviction.cpp rename to tests/cpp/umbp/distributed/test_peer_ssd_eviction.cpp diff --git a/src/umbp/tests/test_peer_ssd_manager.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp similarity index 100% rename from src/umbp/tests/test_peer_ssd_manager.cpp rename to tests/cpp/umbp/distributed/test_peer_ssd_manager.cpp diff --git a/src/umbp/tests/test_peer_ssd_read_rpc.cpp b/tests/cpp/umbp/distributed/test_peer_ssd_read_rpc.cpp similarity index 100% rename from src/umbp/tests/test_peer_ssd_read_rpc.cpp rename to tests/cpp/umbp/distributed/test_peer_ssd_read_rpc.cpp diff --git a/tests/cpp/umbp/distributed/test_route_get_strategy.cpp b/tests/cpp/umbp/distributed/test_route_get_strategy.cpp index 269c91bed..bdea94b93 100644 --- a/tests/cpp/umbp/distributed/test_route_get_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_route_get_strategy.cpp @@ -130,5 +130,67 @@ TEST(CustomRouteGetStrategyTest, FallsBackToFirstWhenNoLocalReplica) { EXPECT_EQ(selected.node_id, "node-a"); } +// ---- BatchSelect tests ---- + +TEST(RouteGetBatchSelectTest, ReturnsOneResultPerKeyInOrder) { + TierPriorityRouteGetStrategy strategy; + std::vector> per_key = { + {MakeLoc("node-a", "loc-1")}, + {MakeLoc("node-b", "loc-2")}, + {MakeLoc("node-c", "loc-3")}, + }; + + auto out = strategy.BatchSelect(per_key, "requester"); + ASSERT_EQ(out.size(), 3u); + EXPECT_EQ(out[0].node_id, "node-a"); + EXPECT_EQ(out[1].node_id, "node-b"); + EXPECT_EQ(out[2].node_id, "node-c"); +} + +TEST(RouteGetBatchSelectTest, EmptyCandidateListLeftAsDefault) { + RandomRouteGetStrategy strategy; + std::vector> per_key = { + {MakeLoc("node-a", "loc-1")}, + {}, // not routed + {MakeLoc("node-c", "loc-3")}, + }; + + auto out = strategy.BatchSelect(per_key, "requester"); + ASSERT_EQ(out.size(), 3u); + EXPECT_EQ(out[0].node_id, "node-a"); + EXPECT_TRUE(out[1].node_id.empty()); // default Location => caller skips + EXPECT_EQ(out[2].node_id, "node-c"); +} + +TEST(RouteGetBatchSelectTest, TierPriorityPicksFastestTierPerKey) { + TierPriorityRouteGetStrategy strategy; + Location hbm = MakeLoc("node-a", "loc-1"); // HBM by default in MakeLoc + Location ssd = MakeLoc("node-b", "loc-2"); + ssd.tier = TierType::SSD; + + std::vector> per_key = {{ssd, hbm}}; + auto out = strategy.BatchSelect(per_key, "requester"); + ASSERT_EQ(out.size(), 1u); + EXPECT_EQ(out[0].tier, TierType::HBM); + EXPECT_EQ(out[0].node_id, "node-a"); +} + +TEST(RouteGetBatchSelectTest, CustomStrategyUsesBaseDefaultLoop) { + // LocalityAwareGetStrategy overrides only Select(); BatchSelect must still + // work via the base-class default that loops Select(). + LocalityAwareGetStrategy strategy; + std::vector> per_key = { + {MakeLoc("node-a", "loc-1"), MakeLoc("node-b", "loc-2")}, + {}, + {MakeLoc("node-a", "loc-3"), MakeLoc("node-c", "loc-4")}, + }; + + auto out = strategy.BatchSelect(per_key, "node-c"); + ASSERT_EQ(out.size(), 3u); + EXPECT_EQ(out[0].node_id, "node-a"); // no local replica -> first + EXPECT_TRUE(out[1].node_id.empty()); // empty -> Select not invoked + EXPECT_EQ(out[2].node_id, "node-c"); // local replica preferred +} + } // namespace } // namespace mori::umbp diff --git a/tests/cpp/umbp/distributed/test_route_put_strategy.cpp b/tests/cpp/umbp/distributed/test_route_put_strategy.cpp index fe298445c..61c81456c 100644 --- a/tests/cpp/umbp/distributed/test_route_put_strategy.cpp +++ b/tests/cpp/umbp/distributed/test_route_put_strategy.cpp @@ -41,8 +41,8 @@ ClientRecord MakeClient(const std::string& node_id, const std::string& addr, rec.node_address = addr; rec.peer_address = addr; rec.status = ClientStatus::ALIVE; - rec.last_heartbeat = std::chrono::steady_clock::now(); - rec.registered_at = std::chrono::steady_clock::now(); + rec.last_heartbeat = std::chrono::system_clock::now(); + rec.registered_at = std::chrono::system_clock::now(); rec.tier_capacities = std::move(caps); return rec; } diff --git a/src/umbp/tests/test_router_dedup.cpp b/tests/cpp/umbp/distributed/test_router_dedup.cpp similarity index 56% rename from src/umbp/tests/test_router_dedup.cpp rename to tests/cpp/umbp/distributed/test_router_dedup.cpp index 893b15820..b871bb211 100644 --- a/src/umbp/tests/test_router_dedup.cpp +++ b/tests/cpp/umbp/distributed/test_router_dedup.cpp @@ -24,14 +24,14 @@ // with already_exists=true and bypass node selection. #include +#include #include #include #include #include #include -#include "umbp/distributed/master/client_registry.h" -#include "umbp/distributed/master/global_block_index.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" #include "umbp/distributed/routing/router.h" #include "umbp/distributed/types.h" @@ -47,19 +47,40 @@ std::map MakeDramCaps(uint64_t total = 8 * kGB) { return caps; } +ClientRegistration MakeRegistration(const std::string& node_id, const std::string& node_address, + const std::string& peer_address) { + ClientRegistration reg; + reg.node_id = node_id; + reg.node_address = node_address; + reg.tier_capacities = MakeDramCaps(); + reg.peer_address = peer_address; + return reg; +} + +// Register `node_id` ALIVE and apply one ADD event for `key` so it has a block +// location in the store. Under the merged store a location can only be created +// through an ApplyHeartbeat from a registered (alive) node — locations no +// longer exist independently of a client record the way the old +// GlobalBlockIndex allowed. +void RegisterWithKey(InMemoryMasterMetadataStore& store, const std::string& node_id, + const std::string& key, std::chrono::system_clock::time_point now) { + ASSERT_TRUE(store.RegisterClient(MakeRegistration(node_id, node_id + ":1", node_id + ":peer"), + now, std::chrono::seconds{30})); + auto hb = store.ApplyHeartbeat(node_id, /*seq=*/1, now, MakeDramCaps(), + {KvEvent{KvEvent::Kind::ADD, key, TierType::DRAM, 4096}}, + /*is_full_sync=*/false); + ASSERT_EQ(hb.status, HeartbeatResult::APPLIED); +} + } // namespace // Indexed keys are marked already_exists; unknown keys still routed. TEST(RouterDedup, BatchRoutePutMarksAlreadyExistsForIndexedKey) { - GlobalBlockIndex index; - ClientRegistry registry(ClientRegistryConfig{}, index); - Router router(index, registry); + const auto now = std::chrono::system_clock::now(); + InMemoryMasterMetadataStore store; + Router router(store); - ASSERT_TRUE(registry.RegisterClient("node-a", "node-a:1", MakeDramCaps(), - /*peer_address=*/"node-a:peer")); - ASSERT_EQ( - index.ApplyEvents("node-a", {KvEvent{KvEvent::Kind::ADD, "key-X", TierType::DRAM, 4096}}), - 1u); + RegisterWithKey(store, "node-a", "key-X", now); std::vector keys{"key-X", "key-Y"}; std::vector sizes{4096, 4096}; @@ -77,22 +98,73 @@ TEST(RouterDedup, BatchRoutePutMarksAlreadyExistsForIndexedKey) { EXPECT_EQ(results[1]->node_id, "node-a"); } +// already_exists wins over an unroutable Put: an existing key is marked +// kAlreadyExists even when no node can accept the write. In the old design +// "no node" meant an empty registry while a foreign node owned the key; under +// the merged store a location can't outlive its alive owner, so the +// unroutable condition is expressed by excluding the only candidate node. The +// property under test is unchanged: dedup wins over node selection. +TEST(RouterDedup, BatchRoutePutAlreadyExistsBypassesUnroutablePut) { + const auto now = std::chrono::system_clock::now(); + InMemoryMasterMetadataStore store; + Router router(store); + + RegisterWithKey(store, "node-a", "key-X", now); + + std::vector keys{"key-X", "key-Y"}; + std::vector sizes{4096, 4096}; + std::unordered_set excludes{"node-a"}; // no routable target left + + auto results = router.BatchRoutePut(keys, "requester", sizes, excludes); + ASSERT_EQ(results.size(), 2u); + + ASSERT_TRUE(results[0].has_value()); + EXPECT_EQ(results[0]->outcome, RoutePutOutcome::kAlreadyExists); + EXPECT_FALSE(results[1].has_value()); // distinct from kAlreadyExists +} + +// Single-key RoutePut delegates to the batch path, so master-side dedup applies: +// an indexed key returns kAlreadyExists while an unknown key still routes. This +// locks in the delegation; without it RoutePut would silently skip dedup. +TEST(RouterDedup, RoutePutMarksAlreadyExistsForIndexedKey) { + const auto now = std::chrono::system_clock::now(); + InMemoryMasterMetadataStore store; + Router router(store); + + RegisterWithKey(store, "node-a", "key-X", now); + + std::unordered_set excludes; + + auto dedup = router.RoutePut("key-X", "requester", 4096, excludes); + ASSERT_TRUE(dedup.has_value()); + EXPECT_EQ(dedup->outcome, RoutePutOutcome::kAlreadyExists); + EXPECT_TRUE(dedup->node_id.empty()); + + auto routed = router.RoutePut("key-Y", "requester", 4096, excludes); + ASSERT_TRUE(routed.has_value()); + EXPECT_EQ(routed->outcome, RoutePutOutcome::kRouted); + EXPECT_EQ(routed->node_id, "node-a"); +} + // Dedup hits never enter the planner, so they consume no projected capacity. -// node-a fits exactly one 6GB block on HBM. key-X is an indexed dedup hit +// node-a fits exactly one 6GB block on DRAM. key-X is an indexed dedup hit // sized 6GB; if dedup had deducted capacity, node-a would have 0 left and the // new key-Y could not route. Since dedup does not deduct, key-Y still routes. TEST(RouterDedup, BatchRoutePutDedupDoesNotConsumeCapacity) { - GlobalBlockIndex index; - ClientRegistry registry(ClientRegistryConfig{}, index); - Router router(index, registry); + const auto now = std::chrono::system_clock::now(); + InMemoryMasterMetadataStore store; + Router router(store); + // node-a holds exactly one 6GB slot on DRAM, already occupied by key-X. std::map caps; - caps[TierType::HBM] = {6 * kGB, 6 * kGB}; - ASSERT_TRUE(registry.RegisterClient("node-a", "node-a:1", caps, - /*peer_address=*/"node-a:peer")); - ASSERT_EQ( - index.ApplyEvents("node-a", {KvEvent{KvEvent::Kind::ADD, "key-X", TierType::HBM, 6 * kGB}}), - 1u); + caps[TierType::DRAM] = {6 * kGB, 6 * kGB}; + ClientRegistration reg = MakeRegistration("node-a", "node-a:1", "node-a:peer"); + reg.tier_capacities = caps; + ASSERT_TRUE(store.RegisterClient(reg, now, std::chrono::seconds{30})); + auto hb = store.ApplyHeartbeat("node-a", /*seq=*/1, now, caps, + {KvEvent{KvEvent::Kind::ADD, "key-X", TierType::DRAM, 6 * kGB}}, + /*is_full_sync=*/false); + ASSERT_EQ(hb.status, HeartbeatResult::APPLIED); std::vector keys{"key-X", "key-Y"}; std::vector sizes{6 * kGB, 6 * kGB}; @@ -107,20 +179,21 @@ TEST(RouterDedup, BatchRoutePutDedupDoesNotConsumeCapacity) { ASSERT_TRUE(results[1].has_value()); EXPECT_EQ(results[1]->outcome, RoutePutOutcome::kRouted); EXPECT_EQ(results[1]->node_id, "node-a"); - EXPECT_EQ(results[1]->tier, TierType::HBM); + EXPECT_EQ(results[1]->tier, TierType::DRAM); } // Projected capacity is batch-local: BatchRoutePut must not write the -// deductions back to the registry's real capacity. +// deductions back to the store's real client capacity. TEST(RouterDedup, BatchRoutePutDoesNotWriteBackProjectedCapacity) { - GlobalBlockIndex index; - ClientRegistry registry(ClientRegistryConfig{}, index); - Router router(index, registry); + const auto now = std::chrono::system_clock::now(); + InMemoryMasterMetadataStore store; + Router router(store); std::map caps; - caps[TierType::HBM] = {80 * kGB, 10 * kGB}; - ASSERT_TRUE(registry.RegisterClient("node-a", "node-a:1", caps, - /*peer_address=*/"node-a:peer")); + caps[TierType::DRAM] = {80 * kGB, 10 * kGB}; + ClientRegistration reg = MakeRegistration("node-a", "node-a:1", "node-a:peer"); + reg.tier_capacities = caps; + ASSERT_TRUE(store.RegisterClient(reg, now, std::chrono::seconds{30})); std::vector keys{"key-A", "key-B"}; std::vector sizes{1 * kGB, 2 * kGB}; @@ -131,59 +204,9 @@ TEST(RouterDedup, BatchRoutePutDoesNotWriteBackProjectedCapacity) { ASSERT_TRUE(results[0].has_value()); EXPECT_EQ(results[0]->outcome, RoutePutOutcome::kRouted); - auto alive = registry.GetAliveClients(); - ASSERT_EQ(alive.size(), 1u); - EXPECT_EQ(alive[0].tier_capacities.at(TierType::HBM).available_bytes, 10 * kGB); -} - -// already_exists wins over no-alive-client: caller drops Put even if -// registry is empty (some other node owns the key). -TEST(RouterDedup, BatchRoutePutAlreadyExistsBypassesNoAliveClient) { - GlobalBlockIndex index; - ClientRegistry registry(ClientRegistryConfig{}, index); - Router router(index, registry); - - ASSERT_EQ( - index.ApplyEvents("node-a", {KvEvent{KvEvent::Kind::ADD, "key-X", TierType::DRAM, 4096}}), - 1u); - - std::vector keys{"key-X", "key-Y"}; - std::vector sizes{4096, 4096}; - std::unordered_set excludes; - - auto results = router.BatchRoutePut(keys, "requester", sizes, excludes); - ASSERT_EQ(results.size(), 2u); - - ASSERT_TRUE(results[0].has_value()); - EXPECT_EQ(results[0]->outcome, RoutePutOutcome::kAlreadyExists); - EXPECT_FALSE(results[1].has_value()); // distinct from kAlreadyExists -} - -// Single-key RoutePut delegates to the batch path, so master-side dedup applies: -// an indexed key returns kAlreadyExists while an unknown key still routes. This -// locks in the delegation; without it RoutePut would silently skip dedup. -TEST(RouterDedup, RoutePutMarksAlreadyExistsForIndexedKey) { - GlobalBlockIndex index; - ClientRegistry registry(ClientRegistryConfig{}, index); - Router router(index, registry); - - ASSERT_TRUE(registry.RegisterClient("node-a", "node-a:1", MakeDramCaps(), - /*peer_address=*/"node-a:peer")); - ASSERT_EQ( - index.ApplyEvents("node-a", {KvEvent{KvEvent::Kind::ADD, "key-X", TierType::DRAM, 4096}}), - 1u); - - std::unordered_set excludes; - - auto dedup = router.RoutePut("key-X", "requester", 4096, excludes); - ASSERT_TRUE(dedup.has_value()); - EXPECT_EQ(dedup->outcome, RoutePutOutcome::kAlreadyExists); - EXPECT_TRUE(dedup->node_id.empty()); - - auto routed = router.RoutePut("key-Y", "requester", 4096, excludes); - ASSERT_TRUE(routed.has_value()); - EXPECT_EQ(routed->outcome, RoutePutOutcome::kRouted); - EXPECT_EQ(routed->node_id, "node-a"); + auto rec = store.GetClient("node-a"); + ASSERT_TRUE(rec.has_value()); + EXPECT_EQ(rec->tier_capacities.at(TierType::DRAM).available_bytes, 10 * kGB); } } // namespace mori::umbp diff --git a/src/umbp/tests/test_ssd_copy_pipeline.cpp b/tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp similarity index 100% rename from src/umbp/tests/test_ssd_copy_pipeline.cpp rename to tests/cpp/umbp/distributed/test_ssd_copy_pipeline.cpp diff --git a/src/umbp/tests/test_ssd_read_lease_gating.cpp b/tests/cpp/umbp/distributed/test_ssd_read_lease_gating.cpp similarity index 100% rename from src/umbp/tests/test_ssd_read_lease_gating.cpp rename to tests/cpp/umbp/distributed/test_ssd_read_lease_gating.cpp diff --git a/src/umbp/tests/test_ssd_reliability.cpp b/tests/cpp/umbp/distributed/test_ssd_reliability.cpp similarity index 78% rename from src/umbp/tests/test_ssd_reliability.cpp rename to tests/cpp/umbp/distributed/test_ssd_reliability.cpp index d1d692567..783a988a9 100644 --- a/src/umbp/tests/test_ssd_reliability.cpp +++ b/tests/cpp/umbp/distributed/test_ssd_reliability.cpp @@ -25,7 +25,7 @@ // * the unified owned-location event source merges DRAM + SSD into one // snapshot/delta (so a heartbeat full-sync ships SSD owned keys too); // * a local SSD eviction's REMOVE SSD event converges the master -// GlobalBlockIndex while leaving the DRAM bucket intact; +// metadata store while leaving the DRAM bucket intact; // * tier-priority RouteGet over the real index picks DRAM, then SSD once the // DRAM replica is removed; // * crash-restart leftover is discarded at startup; @@ -33,18 +33,21 @@ // // (copy-pin vs DRAM evict is covered by test_ssd_copy_pipeline's // EvictBlockedWhilePinnedThenAllowedAfterRelease; seq-gap -> full-sync by -// test_global_block_index_events' ClientRegistryHeartbeat.SeqGap*.) +// test_in_memory_master_metadata_store's heartbeat SeqGap cases.) #include +#include #include +#include #include #include #include #include +#include #include #include -#include "umbp/distributed/master/global_block_index.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" #include "umbp/distributed/peer/owned_location_source.h" #include "umbp/distributed/peer/peer_ssd_manager.h" #include "umbp/distributed/routing/route_get_strategy.h" @@ -138,6 +141,29 @@ int CountTier(const std::vector& events, KvEvent::Kind kind, TierType t return n; } +// Under the merged store a block location can only be created through an +// ApplyHeartbeat from a registered (alive) node — locations no longer exist +// independently of a client record the way the old GlobalBlockIndex allowed. +// These helpers register a node once and apply heartbeat deltas with an +// ascending seq, standing in for the old GlobalBlockIndex::ApplyEvents. +constexpr uint64_t kGB = 1024ULL * 1024 * 1024; + +std::map MakeCaps() { + std::map caps; + caps[TierType::DRAM] = {8 * kGB, 8 * kGB}; + caps[TierType::SSD] = {8 * kGB, 8 * kGB}; + return caps; +} + +ClientRegistration MakeReg(const std::string& node_id) { + ClientRegistration reg; + reg.node_id = node_id; + reg.node_address = node_id + ":1"; + reg.peer_address = node_id + ":peer"; + reg.tier_capacities = MakeCaps(); + return reg; +} + // A canned owned-location source standing in for PeerDramAllocator so the // aggregation can be tested without standing up a DRAM allocator. class FakeOwnedSource : public OwnedLocationSource { @@ -198,58 +224,86 @@ TEST(SsdReliability, DeltaDrainMergesDramAndSsdEvents) { } // --------------------------------------------------------------------------- -// SSD local eviction -> REMOVE SSD -> master GlobalBlockIndex converges. +// SSD local eviction -> REMOVE SSD -> master metadata store converges. // --------------------------------------------------------------------------- // A key mirrored on DRAM + SSD of one owner: a local SSD eviction emits -// REMOVE SSD, and applying that to the master index drops only the SSD bucket +// REMOVE SSD, and applying that to the master store drops only the SSD bucket // (the DRAM replica, owned independently, stays routable). TEST(SsdReliability, LocalSsdEvictionRemoveConvergesMasterIndex) { - GlobalBlockIndex idx; + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); + ASSERT_TRUE(store.RegisterClient(MakeReg("owner"), now, std::chrono::seconds{30})); auto be = std::make_unique(1'000'000); PeerSsdManager ssd(std::move(be), 0.9, 0.7); // DRAM replica added independently (a DRAM owner would emit this). - idx.ApplyEvents("owner", {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}}); - // SSD copy lands -> ADD SSD drained into the index. + ASSERT_EQ(store + .ApplyHeartbeat("owner", /*seq=*/1, now, MakeCaps(), + {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}}, + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); + // SSD copy lands -> ADD SSD drained into the store. ASSERT_TRUE(ssd.Write("k", OneSeg(std::string(100, 'x')), 100)); - idx.ApplyEvents("owner", ssd.DrainPendingEvents()); + ASSERT_EQ(store + .ApplyHeartbeat("owner", /*seq=*/2, now, MakeCaps(), ssd.DrainPendingEvents(), + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); - auto both = idx.Lookup("k"); + auto both = store.LookupBlock("k"); ASSERT_TRUE(HasLoc(both, "owner", TierType::DRAM)); ASSERT_TRUE(HasLoc(both, "owner", TierType::SSD)); - // Local SSD eviction -> REMOVE SSD -> index drops only the SSD bucket. + // Local SSD eviction -> REMOVE SSD -> store drops only the SSD bucket. ASSERT_TRUE(ssd.Evict("k")); auto ssd_events = ssd.DrainPendingEvents(); EXPECT_EQ(CountTier(ssd_events, KvEvent::Kind::REMOVE, TierType::SSD), 1); - idx.ApplyEvents("owner", ssd_events); + ASSERT_EQ(store + .ApplyHeartbeat("owner", /*seq=*/3, now, MakeCaps(), ssd_events, + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); - auto after = idx.Lookup("k"); + auto after = store.LookupBlock("k"); EXPECT_TRUE(HasLoc(after, "owner", TierType::DRAM)); // DRAM replica still routable EXPECT_FALSE(HasLoc(after, "owner", TierType::SSD)); // SSD bucket converged away } // --------------------------------------------------------------------------- -// Tier-priority RouteGet over the real index: DRAM first, SSD after evict. +// Tier-priority RouteGet over the real store: DRAM first, SSD after evict. // --------------------------------------------------------------------------- TEST(SsdReliability, TierPriorityRoutesDramThenSsdAfterDramRemoved) { - GlobalBlockIndex idx; - idx.ApplyEvents("owner", {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}, - KvEvent{KvEvent::Kind::ADD, "k", TierType::SSD, 100}}); + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); + const std::unordered_set kNoExclude; + ASSERT_TRUE(store.RegisterClient(MakeReg("owner"), now, std::chrono::seconds{30})); + ASSERT_EQ(store + .ApplyHeartbeat("owner", /*seq=*/1, now, MakeCaps(), + {KvEvent{KvEvent::Kind::ADD, "k", TierType::DRAM, 100}, + KvEvent{KvEvent::Kind::ADD, "k", TierType::SSD, 100}}, + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); TierPriorityRouteGetStrategy strategy; - auto locs = idx.BatchLookupForRouteGet({"k"}, {}, std::chrono::seconds{10}); + auto locs = store.BatchLookupBlockForRouteGet({"k"}, kNoExclude, now, std::chrono::seconds{10}); ASSERT_EQ(locs.size(), 1u); auto dram_pick = strategy.Select(locs[0], "reader"); EXPECT_EQ(dram_pick.tier, TierType::DRAM) << "prefers the fast DRAM replica"; // DRAM evicted -> only the SSD bucket remains -> RouteGet must serve from SSD. - idx.ApplyEvents("owner", {KvEvent{KvEvent::Kind::REMOVE, "k", TierType::DRAM, 0}}); - auto locs2 = idx.BatchLookupForRouteGet({"k"}, {}, std::chrono::seconds{10}); + ASSERT_EQ(store + .ApplyHeartbeat("owner", /*seq=*/2, now, MakeCaps(), + {KvEvent{KvEvent::Kind::REMOVE, "k", TierType::DRAM, 0}}, + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); + auto locs2 = store.BatchLookupBlockForRouteGet({"k"}, kNoExclude, now, std::chrono::seconds{10}); ASSERT_EQ(locs2.size(), 1u); auto ssd_pick = strategy.Select(locs2[0], "reader"); EXPECT_EQ(ssd_pick.tier, TierType::SSD) << "falls back to the surviving SSD replica"; diff --git a/src/umbp/tests/test_tier_priority_route_get.cpp b/tests/cpp/umbp/distributed/test_tier_priority_route_get.cpp similarity index 100% rename from src/umbp/tests/test_tier_priority_route_get.cpp rename to tests/cpp/umbp/distributed/test_tier_priority_route_get.cpp diff --git a/tests/cpp/umbp/distributed/test_umbp_tags.cpp b/tests/cpp/umbp/distributed/test_umbp_tags.cpp index 51b965198..3a0084c32 100644 --- a/tests/cpp/umbp/distributed/test_umbp_tags.cpp +++ b/tests/cpp/umbp/distributed/test_umbp_tags.cpp @@ -21,9 +21,9 @@ // SOFTWARE. // Tests for the client-tag feature: // -// Suite 1 — ClientRegistryTagsTest -// Unit tests directly on ClientRegistry: verify tags are stored on -// RegisterClient and returned verbatim by GetClientTags. +// Suite 1 — ClientTagsTest +// Unit tests directly on InMemoryMasterMetadataStore: verify tags are stored +// on RegisterClient and returned verbatim by GetClientTags. // // Suite 2 — MasterClientTagsE2ETest // Integration test: MasterClient registers with tags via gRPC, the real @@ -53,7 +53,7 @@ #include "umbp.grpc.pb.h" #include "umbp/distributed/config.h" -#include "umbp/distributed/master/client_registry.h" +#include "umbp/distributed/master/in_memory_master_metadata_store.h" #include "umbp/distributed/master/master_client.h" #include "umbp/distributed/master/master_server.h" #include "umbp/distributed/types.h" @@ -94,54 +94,77 @@ static bool WaitFor(std::function pred, std::chrono::milliseconds timeou } // --------------------------------------------------------------------------- -// Suite 1: ClientRegistry unit tests (no gRPC) +// Suite 1: InMemoryMasterMetadataStore unit tests (no gRPC) // --------------------------------------------------------------------------- -TEST(ClientRegistryTagsTest, TagsStoredOnRegister) { - ClientRegistry reg(ClientRegistryConfig{}); +// Build a registration for `node_id` carrying `tags`. The merged store creates +// client records through RegisterClient(ClientRegistration, now, stale_after), +// replacing the old ClientRegistry::RegisterClient(node, addr, caps, ...) form. +ClientRegistration MakeReg(const std::string& node_id, const std::string& node_address, + std::vector tags = {}) { + ClientRegistration reg; + reg.node_id = node_id; + reg.node_address = node_address; + reg.tags = std::move(tags); + return reg; +} + +TEST(ClientTagsTest, TagsStoredOnRegister) { + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); const std::vector tags = {"sgl_role=prefill", "env=test"}; - ASSERT_TRUE(reg.RegisterClient("n1", "127.0.0.1:9001", {}, /*peer=*/"", - /*engine=*/{}, tags)); - EXPECT_EQ(reg.GetClientTags("n1"), tags); + ASSERT_TRUE( + store.RegisterClient(MakeReg("n1", "127.0.0.1:9001", tags), now, std::chrono::seconds{30})); + EXPECT_EQ(store.GetClientTags("n1"), tags); } -TEST(ClientRegistryTagsTest, EmptyTagsReturnedForUnknownNode) { - ClientRegistry reg(ClientRegistryConfig{}); - EXPECT_TRUE(reg.GetClientTags("ghost").empty()); +TEST(ClientTagsTest, EmptyTagsReturnedForUnknownNode) { + InMemoryMasterMetadataStore store; + EXPECT_TRUE(store.GetClientTags("ghost").empty()); } -TEST(ClientRegistryTagsTest, EmptyTagsWhenNoneProvided) { - ClientRegistry reg(ClientRegistryConfig{}); - ASSERT_TRUE(reg.RegisterClient("n1", "127.0.0.1:9002", {})); - EXPECT_TRUE(reg.GetClientTags("n1").empty()); +TEST(ClientTagsTest, EmptyTagsWhenNoneProvided) { + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); + ASSERT_TRUE(store.RegisterClient(MakeReg("n1", "127.0.0.1:9002"), now, std::chrono::seconds{30})); + EXPECT_TRUE(store.GetClientTags("n1").empty()); } -TEST(ClientRegistryTagsTest, TagsUnchangedByHeartbeat) { - ClientRegistry reg(ClientRegistryConfig{}); +TEST(ClientTagsTest, TagsUnchangedByHeartbeat) { + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); const std::vector tags = {"sgl_role=decode"}; - ASSERT_TRUE(reg.RegisterClient("n1", "127.0.0.1:9003", {}, "", {}, tags)); + ASSERT_TRUE( + store.RegisterClient(MakeReg("n1", "127.0.0.1:9003", tags), now, std::chrono::seconds{30})); - uint64_t acked = 0; - bool request_full_sync = false; - reg.Heartbeat("n1", {}, {}, /*is_full_sync=*/false, 0, &acked, &request_full_sync); + ASSERT_EQ(store + .ApplyHeartbeat("n1", /*seq=*/1, now, /*caps=*/{}, /*events=*/{}, + /*is_full_sync=*/false) + .status, + HeartbeatResult::APPLIED); - EXPECT_EQ(reg.GetClientTags("n1"), tags); + EXPECT_EQ(store.GetClientTags("n1"), tags); } -TEST(ClientRegistryTagsTest, TagsClearedAfterUnregister) { - ClientRegistry reg(ClientRegistryConfig{}); - ASSERT_TRUE(reg.RegisterClient("n1", "127.0.0.1:9004", {}, "", {}, {"sgl_role=prefill"})); - reg.UnregisterClient("n1"); - EXPECT_TRUE(reg.GetClientTags("n1").empty()); +TEST(ClientTagsTest, TagsClearedAfterUnregister) { + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); + ASSERT_TRUE(store.RegisterClient(MakeReg("n1", "127.0.0.1:9004", {"sgl_role=prefill"}), now, + std::chrono::seconds{30})); + store.UnregisterClient("n1"); + EXPECT_TRUE(store.GetClientTags("n1").empty()); } -TEST(ClientRegistryTagsTest, MultipleNodesHaveIndependentTags) { - ClientRegistry reg(ClientRegistryConfig{}); - ASSERT_TRUE(reg.RegisterClient("p", "127.0.0.1:9005", {}, "", {}, {"sgl_role=prefill"})); - ASSERT_TRUE(reg.RegisterClient("d", "127.0.0.1:9006", {}, "", {}, {"sgl_role=decode"})); +TEST(ClientTagsTest, MultipleNodesHaveIndependentTags) { + InMemoryMasterMetadataStore store; + const auto now = std::chrono::system_clock::now(); + ASSERT_TRUE(store.RegisterClient(MakeReg("p", "127.0.0.1:9005", {"sgl_role=prefill"}), now, + std::chrono::seconds{30})); + ASSERT_TRUE(store.RegisterClient(MakeReg("d", "127.0.0.1:9006", {"sgl_role=decode"}), now, + std::chrono::seconds{30})); - EXPECT_EQ(reg.GetClientTags("p"), std::vector{"sgl_role=prefill"}); - EXPECT_EQ(reg.GetClientTags("d"), std::vector{"sgl_role=decode"}); + EXPECT_EQ(store.GetClientTags("p"), std::vector{"sgl_role=prefill"}); + EXPECT_EQ(store.GetClientTags("d"), std::vector{"sgl_role=decode"}); } // --------------------------------------------------------------------------- @@ -282,7 +305,7 @@ TEST_F(MasterClientTagsE2ETest, EmptyTagsSentWhenNoneConfigured) { // Verify that MasterClient::AddCounter forwards labels to the server and // that a capturing server would see the node label. The real tag-injection // into ReportMetrics base labels is exercised in the real-MasterServer suite -// below because CapturingMasterService doesn't run ClientRegistry. +// below because CapturingMasterService doesn't run the master metadata store. TEST_F(MasterClientTagsE2ETest, ReportMetricsCarriesNodeId) { setenv("UMBP_METRICS_REPORT_INTERVAL_MS", "50", 1); client_ = std::make_unique(MakeConfig({"sgl_role=decode"}));