feat(umbp): converge eviction onto one path; nodes evict locally by default

src/umbp/distributed/master/evict_strategy.cpp

@@ -26,7 +26,7 @@
 
 namespace mori::umbp {
 
-std::unordered_map<std::string, std::vector<std::string>> LruMasterEvictStrategy::SelectVictims(
+std::unordered_map<std::string, std::vector<EvictionVictim>> LruMasterEvictStrategy::SelectVictims(
     std::vector<EvictionCandidate> candidates,
     std::unordered_map<std::string, std::map<TierType, int64_t>> bytes_to_free) {
   // Oldest-access first (LRU); no tiebreak (peers don't ship depth in KvEvent).
@@ -36,16 +36,17 @@ std::unordered_map<std::string, std::vector<std::string>> LruMasterEvictStrategy
             });
 
   // Walk oldest-first, taking a candidate while its (node, tier) has budget
-  // left.  Group by node so each peer gets a single EvictKey keys[].
-  std::unordered_map<std::string, std::vector<std::string>> per_node_keys;
+  // left.  Group by node so each peer gets a single EvictKey.  The victim keeps
+  // the tier it was charged to: that is the medium the peer must free.
+  std::unordered_map<std::string, std::vector<EvictionVictim>> per_node_victims;
   for (const auto& c : candidates) {
     auto& tier_budget = bytes_to_free[c.location.node_id];
     auto it = tier_budget.find(c.location.tier);
     if (it == tier_budget.end() || it->second <= 0) continue;
-    per_node_keys[c.location.node_id].push_back(c.key);
+    per_node_victims[c.location.node_id].push_back(EvictionVictim{c.key, c.location.tier});
     it->second -= static_cast<int64_t>(c.size);
   }
-  return per_node_keys;
+  return per_node_victims;
 }
 
 }  // namespace mori::umbp

src/umbp/distributed/master/eviction_manager.cpp

@@ -21,6 +21,7 @@
 // SOFTWARE.
 #include "umbp/distributed/master/eviction_manager.h"
 
+#include <algorithm>
 #include <chrono>
 #include <cstdint>
 #include <iterator>
@@ -49,12 +50,59 @@ EvictionManager::EvictionManager(IMasterMetadataStore& store, const EvictionConf
 
 EvictionManager::~EvictionManager() { Stop(); }
 
+void EvictionManager::SetPeerDriven(const std::string& node_id, bool peer_driven) {
+  std::lock_guard<std::mutex> lock(peer_driven_mutex_);
+  if (peer_driven) {
+    if (peer_driven_.insert(node_id).second) {
+      MORI_UMBP_INFO("[EvictionManager] node {} evicts for itself; no EvictKey rounds", node_id);
+    }
+  } else {
+    peer_driven_.erase(node_id);
+  }
+}
+
+bool EvictionManager::PeerDriven(const std::string& node_id) const {
+  std::lock_guard<std::mutex> lock(peer_driven_mutex_);
+  return peer_driven_.count(node_id) != 0;
+}
+
+std::vector<EvictionCandidate> SelectPeerEvictionCandidates(
+    const IMasterMetadataStore& store, const std::string& node_id, TierType tier,
+    const std::string& logical_tier, uint64_t bytes, size_t max_keys,
+    std::chrono::system_clock::time_point now) {
+  if (max_keys == 0 || max_keys > kMaxPeerEvictionCandidates) {
+    max_keys = kMaxPeerEvictionCandidates;
+  }
+  const NodeTierKey bucket{node_id, tier};
+  // A logical-tier filter drops rows after the store returns them, so it
+  // cannot cap.
+  auto by_bucket = store.EnumerateEvictionCandidates(
+      {bucket}, EvictionOrder::kLeastRecentlyAccessed, logical_tier.empty() ? max_keys : 0, now);
+  auto it = by_bucket.find(bucket);
+  if (it == by_bucket.end()) return {};
+  auto& rows = it->second;
+  // The order is only a hint to the store.
+  std::sort(rows.begin(), rows.end(),
+            [](const auto& a, const auto& b) { return a.last_accessed_at < b.last_accessed_at; });
+  std::vector<EvictionCandidate> out;
+  uint64_t covered = 0;
+  for (auto& row : rows) {
+    if (out.size() >= max_keys || covered >= bytes) break;
+    if (!logical_tier.empty() && row.location.logical_tier != logical_tier) continue;
+    covered += row.size;
+    out.push_back(std::move(row));
+  }
+  return out;
+}
+
 void EvictionManager::Start() {
   if (running_.load(std::memory_order_relaxed)) return;
   running_.store(true, std::memory_order_relaxed);
   thread_ = std::thread(&EvictionManager::EvictionLoop, this);
-  MORI_UMBP_INFO("[EvictionManager] Started (interval={}s, high={}, low={})",
-                 config_.check_interval.count(), config_.high_watermark, config_.low_watermark);
+  MORI_UMBP_INFO(
+      "[EvictionManager] Started (interval={}s, high={}, low={}, max_bytes_per_round={})",
+      config_.check_interval.count(), config_.high_watermark, config_.low_watermark,
+      config_.max_evict_bytes_per_round);
 }
 
 void EvictionManager::Stop() {
@@ -97,6 +145,7 @@ void EvictionManager::RunOnce() {
   std::map<NodeTierKey, uint64_t> bytes_to_free;
 
   for (const auto& client : clients) {
+    if (PeerDriven(client.node_id)) continue;
     for (const auto& [tier, cap] : client.tier_capacities) {
       // No tier is special-cased here any more (backend-agnostic refactor
       // Phase 4).  The old `if (tier == SSD) continue` existed because EvictKey
@@ -118,28 +167,39 @@ void EvictionManager::RunOnce() {
   }
 
   if (bytes_to_free.empty()) return;
-  MORI_UMBP_INFO("[EvictionManager] {} overloaded node-tiers detected", bytes_to_free.size());
+
+  // A round frees each bucket's overage in bytes, optionally capped per round.
+  if (config_.max_evict_bytes_per_round > 0) {
+    for (auto& [ntk, bytes] : bytes_to_free) {
+      bytes = std::min(bytes, config_.max_evict_bytes_per_round);
+    }
+  }
+  uint64_t budget_total = 0;
+  for (const auto& [ntk, bytes] : bytes_to_free) budget_total += bytes;
+  MORI_UMBP_INFO("[EvictionManager] {} overloaded node-tiers detected, {} bytes to free",
+                 bytes_to_free.size(), budget_total);
 
   // 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. Limit each bucket to the configured
-  // round size; a still-overloaded bucket is revisited on the next interval.
+  // decision and never sees the byte budget; the strategy trims to it, so
+  // there is no row cap.
   std::vector<NodeTierKey> 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, config_.evict_batch_size,
-      std::chrono::system_clock::now());
+  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;
   }
 
   // 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[].
+  // grouped by node so each peer gets a single EvictKey.  Victims keep their
+  // tier, so the peer frees the copy that is over budget.
   std::vector<EvictionCandidate> candidates;
   for (auto& [bucket, rows] : candidates_by_bucket) {
     for (auto& c : rows) candidates.push_back(std::move(c));
@@ -149,32 +209,33 @@ void EvictionManager::RunOnce() {
     strategy_budget[ntk.node_id][ntk.tier] = static_cast<int64_t>(bytes);
   }
 
-  auto per_node_keys = strategy_->SelectVictims(std::move(candidates), std::move(strategy_budget));
+  auto per_node_victims =
+      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();
+  for (const auto& [node_id, victims] : per_node_victims) selected += victims.size();
 
   if (selected == 0) return;
 
   MORI_UMBP_INFO("[EvictionManager] Selected {} victims across {} nodes", selected,
-                 per_node_keys.size());
+                 per_node_victims.size());
 
   // Look up peer addresses once per dispatch round.  ClientRegistry
   // owns the (node_id -> peer_address) mapping; we can't ship an
   // EvictKey to a node that has dropped out.
   std::unordered_map<std::string, std::string> node_to_peer;
   for (const auto& client : clients) node_to_peer[client.node_id] = client.peer_address;
 
-  for (auto& [node_id, keys] : per_node_keys) {
+  for (auto& [node_id, victims] : per_node_victims) {
     auto it = node_to_peer.find(node_id);
     if (it == node_to_peer.end() || it->second.empty()) {
       MORI_UMBP_WARN("[EvictionManager] no peer_address for node={} — skipping {} keys", node_id,
-                     keys.size());
+                     victims.size());
       continue;
     }
     if (dispatcher_ == nullptr) {
       MORI_UMBP_DEBUG("[EvictionManager] dispatcher unset; would EvictKey on node={} ({} keys)",
-                      node_id, keys.size());
+                      node_id, victims.size());
       continue;
     }
     // Chunked to the shared gRPC message limit.  An eviction round frees a
@@ -189,19 +250,23 @@ void EvictionManager::RunOnce() {
     // next heartbeat are what shrink the index), and a chunk that fails is
     // retried next round.  So N chunks reach the same state as one message
     // would have, and a partial round still makes progress.
+    // The one-byte tier per key fits in GrpcMaxItemsPerBatch's item overhead.
+    std::vector<std::string> keys;
+    keys.reserve(victims.size());
+    for (const auto& victim : victims) keys.push_back(victim.key);
     size_t sent = 0;
     size_t chunks = 0;
-    while (sent < keys.size()) {
+    while (sent < victims.size()) {
       const size_t take = GrpcMaxItemsPerBatch(keys, sent);
-      std::vector<std::string> chunk(std::make_move_iterator(keys.begin() + sent),
-                                     std::make_move_iterator(keys.begin() + sent + take));
+      std::vector<EvictionVictim> chunk(std::make_move_iterator(victims.begin() + sent),
+                                        std::make_move_iterator(victims.begin() + sent + take));
       sent += take;
       ++chunks;
       dispatcher_->DispatchEvictKey(node_id, it->second, std::move(chunk));
     }
     if (chunks > 1) {
       MORI_UMBP_DEBUG("[EvictionManager] node={} {} keys dispatched in {} chunks", node_id,
-                      keys.size(), chunks);
+                      victims.size(), chunks);
     }
   }
 }

src/umbp/distributed/master/in_memory_master_metadata_store.cpp

@@ -546,6 +546,23 @@ std::vector<std::vector<Location>> InMemoryMasterMetadataStore::BatchLookupBlock
   return out;
 }
 
+void InMemoryMasterMetadataStore::RecordLocalReads(const std::vector<std::string>& keys,
+                                                   std::chrono::system_clock::time_point now) {
+  if (keys.empty()) return;
+  // An access bump is atomic, so the shard's shared lock is enough.
+  std::vector<std::vector<size_t>> 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]);
+      if (it != sh.entries.end()) it->second.RecordAccessAtomic(now);
+    }
+  }
+}
+
 std::vector<bool> InMemoryMasterMetadataStore::BatchExistsBlock(
     const std::vector<std::string>& keys) const {
   std::vector<bool> results(keys.size(), false);

src/umbp/distributed/master/master_client.cpp

@@ -204,6 +204,7 @@ grpc::Status MasterClient::RegisterSelf(
   FillTierCapacities(req.mutable_tier_capacities(), tier_capacities);
   FillLogicalTierCapacities(req.mutable_logical_tier_capacities(), logical_tier_capacities);
   for (const auto& tag : config_.tags) req.add_tags(tag);
+  req.set_peer_driven_eviction(peer_driven_requested_);
 
   ::umbp::RegisterClientResponse resp;
   grpc::ClientContext ctx;
@@ -217,6 +218,7 @@ grpc::Status MasterClient::RegisterSelf(
   if (resp.heartbeat_interval_ms() > 0) heartbeat_interval_ms_ = resp.heartbeat_interval_ms();
   supports_max_allocatable_capacity_ = resp.supports_max_allocatable_bytes();
   supports_logical_tiers_ = resp.supports_logical_tiers();
+  supports_peer_driven_eviction_ = resp.supports_peer_driven_eviction();
   aggregate_backend_capacities_ =
       aggregate_backend_capacities_requested_ && supports_max_allocatable_capacity_;
   {
@@ -445,6 +447,57 @@ grpc::Status MasterClient::BatchLookup(const std::vector<std::string>& keys,
   return grpc::Status::OK;
 }
 
+std::optional<std::vector<EvictionOffer>> MasterClient::GetEvictionCandidates(
+    TierType tier, const std::string& logical_tier, uint64_t bytes, uint32_t max_keys) {
+  ScopedRpcTimer _rpc_timer(this, "GetEvictionCandidates");
+  ::umbp::GetEvictionCandidatesRequest req;
+  req.set_node_id(config_.node_id);
+  req.set_tier(ToProtoTier(tier));
+  req.set_logical_tier(logical_tier);
+  req.set_bytes(bytes);
+  req.set_max_keys(max_keys);
+
+  ::umbp::GetEvictionCandidatesResponse resp;
+  grpc::ClientContext ctx;
+  // Eviction does not wait on this (PeerPool::RankedCandidates); past the
+  // age a peer still uses, an answer is wasted work.
+  ctx.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(3));
+  auto status = GetStub(stub_.get())->GetEvictionCandidates(&ctx, req, &resp);
+  _rpc_timer.SetStatus(status);
+  if (!status.ok()) {
+    static std::atomic<uint64_t> failures{0};
+    const uint64_t n = failures.fetch_add(1, std::memory_order_relaxed) + 1;
+    if (n <= 8 || n % 1000 == 0) {
+      MORI_UMBP_WARN("[Client] GetEvictionCandidates failed ({} so far): {}", n,
+                     status.error_message());
+    }
+    return std::nullopt;
+  }
+  const int n = std::min(resp.keys_size(), resp.sizes_size());
+  std::vector<EvictionOffer> out;
+  out.reserve(static_cast<size_t>(n));
+  for (int i = 0; i < n; ++i) out.push_back(EvictionOffer{resp.keys(i), resp.sizes(i)});
+  return out;
+}
+
+void MasterClient::NoteLocalReads(const std::vector<std::string>& keys) {
+  // Past kFlushAt distinct keys, ship them now rather than wait out the
+  // interval; past kMax (master unreachable) the rest are dropped.
+  constexpr size_t kFlushAt = 4096;
+  constexpr size_t kMax = 65536;
+  if (keys.empty() || !ReportsLocalReads()) return;
+  bool flush = false;
+  {
+    std::lock_guard lock(local_reads_mutex_);
+    for (const auto& key : keys) {
+      if (local_reads_.size() >= kMax) break;
+      local_reads_.insert(key);
+    }
+    flush = local_reads_.size() >= kFlushAt;
+  }
+  if (flush) FlushHeartbeat();
+}
+
 namespace {
 // Auto-flush the heartbeat once this many KvEvents accumulate at the peer, so a
 // completed batch of puts becomes visible at the master without the application
@@ -742,6 +795,7 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::map<TierType, TierCapacit
                                 peer_pool_->LogicalTierCapacities());
     }
     for (const auto& tag : config_.tags) re_req.add_tags(tag);
+    re_req.set_peer_driven_eviction(peer_driven_requested_);
     ::umbp::RegisterClientResponse re_resp;
     grpc::ClientContext re_ctx;
     grpc::Status re_status;
@@ -754,6 +808,7 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::map<TierType, TierCapacit
       supports_max_allocatable_capacity_ =
           re_status.ok() && re_resp.supports_max_allocatable_bytes();
       supports_logical_tiers_ = re_status.ok() && re_resp.supports_logical_tiers();
+      supports_peer_driven_eviction_ = re_status.ok() && re_resp.supports_peer_driven_eviction();
       aggregate_backend_capacities_ =
           aggregate_backend_capacities_requested_ && supports_max_allocatable_capacity_;
       registered_ = true;
@@ -780,6 +835,14 @@ bool MasterClient::SendDeltaHeartbeatLocked(const std::map<TierType, TierCapacit
 
 grpc::Status MasterClient::SendHeartbeatRpcLocked(::umbp::HeartbeatRequest& req,
                                                   ::umbp::HeartbeatResponse* resp) {
+  req.set_peer_driven_eviction(peer_driven_requested_);
+  std::unordered_set<std::string> reads;
+  {
+    std::lock_guard lock(local_reads_mutex_);
+    reads.swap(local_reads_);
+  }
+  req.mutable_local_reads()->Reserve(static_cast<int>(reads.size()));
+  while (!reads.empty()) req.add_local_reads(std::move(reads.extract(reads.begin()).value()));
   grpc::ClientContext ctx;
   ctx.set_deadline(std::chrono::system_clock::now() +
                    std::chrono::milliseconds(HeartbeatRpcTimeoutMs()));