fix(umbp): a remote read checks after the copy that it read the right key

src/umbp/distributed/peer/backend/page_backend.cpp

@@ -491,6 +491,7 @@ bool PageBackend::CommitLocked(uint64_t slot_id, const std::string& key,
   owned.pages = std::move(it->second.pages);
   owned.size = it->second.size;
   owned.hold = std::make_shared<char>();
+  owned.generation = NextCopyGeneration();
   QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, tier_, owned.size});
   owned_[key] = std::move(owned);
   TouchLru(key, owned_[key]);
@@ -546,6 +547,7 @@ ResolvedEntry PageBackend::Resolve(const std::string& key) {
   r.pages = it->second.pages;
   r.size = it->second.size;
   r.page_size = page_size_;
+  r.generation = it->second.generation;
   // Extend the read lease so concurrent Evict reports bytes_freed=0 for
   // this key.  CAS-max rather than a plain store: with a shared lock two
   // readers can renew the same key at once, and neither may shorten it.
@@ -581,6 +583,7 @@ std::vector<ResolvedEntry> PageBackend::BatchResolve(const std::vector<std::stri
     entry.pages = it->second.pages;
     entry.size = it->second.size;
     entry.page_size = page_size_;
+    entry.generation = it->second.generation;
     if (include_descs) entry.descs = BuildBufferDescsLocked(it->second.pages);
     if (reader == ResolveReader::kLocal) {
       entry.hold = it->second.hold;
@@ -630,6 +633,21 @@ std::vector<bool> PageBackend::BatchContains(const std::vector<std::string>& key
   return out;
 }
 
+std::vector<bool> PageBackend::BatchStillHolds(const std::vector<std::string>& keys,
+                                               const std::vector<uint64_t>& generations) const {
+  std::vector<bool> out(keys.size(), false);
+  std::shared_lock<std::shared_mutex> lock(mutex_);
+  for (size_t i = 0; i < keys.size() && i < generations.size(); ++i) {
+    if (generations[i] == 0) {
+      out[i] = true;
+      continue;
+    }
+    const auto it = owned_.find(keys[i]);
+    out[i] = it != owned_.end() && it->second.generation == generations[i];
+  }
+  return out;
+}
+
 // ---------------------------------------------------------------------------
 //  Eviction (skips leased / copy-pinned keys; emits REMOVE)
 // ---------------------------------------------------------------------------

src/umbp/distributed/peer/backend/ssd_backend.cpp

@@ -494,6 +494,7 @@ std::vector<ResolvedEntry> SsdBackend::BatchResolve(const std::vector<std::strin
     }
     out.size = lease.size;
     out.page_size = cfg_.page_size;
+    out.generation = lease.generation;
     if (include_descs && !buffer_desc_.empty()) {
       out.descs = {BufferMemoryDescBytes{kStagingBufferIndex, buffer_desc_}};
     }
@@ -656,8 +657,8 @@ std::vector<ResolvedEntry> SsdBackend::BatchResolve(const std::vector<std::strin
       const auto lease_deadline = local ? std::chrono::steady_clock::time_point{}
                                         : std::chrono::steady_clock::now() + cfg_.read_lease_ttl;
       auto [it, inserted] = read_leases_.emplace(
-          keys[primary],
-          ReadLease{spans[j], outcomes[j].size, lease_deadline, std::make_shared<char>()});
+          keys[primary], ReadLease{spans[j], outcomes[j].size, lease_deadline,
+                                   std::make_shared<char>(), NextCopyGeneration()});
       if (!inserted) {
         ReleaseStagingSpanLocked(spans[j]);
         if (!local) it->second.expires_at = lease_deadline;
@@ -728,6 +729,26 @@ bool SsdBackend::Contains(const std::string& key) const {
   return ssd_ != nullptr && ssd_->Exists(key);
 }
 
+// A remote reader copies out of a staging span, not the drive, so the copy is
+// good only while the lease it resolved is still the one in the table.  Any
+// reclaim drops the lease, and a re-stage of the same key draws a new
+// generation, so an equal generation means the span was never handed out again.
+std::vector<bool> SsdBackend::BatchStillHolds(const std::vector<std::string>& keys,
+                                              const std::vector<uint64_t>& generations) const {
+  std::vector<bool> out(keys.size(), false);
+  std::lock_guard<std::mutex> lk(mutex_);
+  for (size_t i = 0; i < keys.size(); ++i) {
+    const uint64_t gen = i < generations.size() ? generations[i] : 0;
+    if (gen == 0) {
+      out[i] = true;
+      continue;
+    }
+    auto it = read_leases_.find(keys[i]);
+    out[i] = it != read_leases_.end() && it->second.generation == gen;
+  }
+  return out;
+}
+
 std::vector<EvictResult> SsdBackend::Evict(const std::vector<std::string>& keys) {
   std::vector<EvictResult> results;
   results.reserve(keys.size());

src/umbp/distributed/peer/peer_service.cpp

@@ -399,6 +399,7 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se
       entries[i].backend_id = result.backend_id;
       entries[i].size = r.size;
       entries[i].pages = std::move(r.pages);
+      entries[i].generation = r.generation;
       total_bytes += r.size;
 
       if (omit_descs) continue;
@@ -420,6 +421,30 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se
     return grpc::Status::OK;
   }
 
+  // A remote reader's read lease is only a timer: once it lapses, the pages it
+  // resolved can be freed and refilled while its RDMA read still runs.  It asks
+  // afterwards: is each copy it read still the one this peer resolved?
+  // Read-only; it neither takes nor extends a lease.
+  grpc::Status ValidateReads(grpc::ServerContext* /*ctx*/,
+                             const ::umbp::ValidateReadsRequest* request,
+                             ::umbp::ValidateReadsResponse* response) override {
+    const int n = request->keys_size();
+    if (request->backend_id_size() != n || request->generation_size() != n) {
+      return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
+                          "ValidateReads: keys, backend_id and generation differ in length");
+    }
+    std::vector<std::string> keys(request->keys().begin(), request->keys().end());
+    std::vector<uint32_t> backend_ids(request->backend_id().begin(), request->backend_id().end());
+    std::vector<uint64_t> generations(request->generation().begin(), request->generation().end());
+    auto valid = pool_ == nullptr ? std::vector<bool>(keys.size(), false)
+                                  : pool_->ValidateResolved(backend_ids, keys, generations);
+    response->mutable_still_valid()->Reserve(n);
+    for (int i = 0; i < n; ++i) {
+      response->add_still_valid(static_cast<size_t>(i) < valid.size() && valid[i]);
+    }
+    return grpc::Status::OK;
+  }
+
  private:
   MediumBackend* Backend(uint32_t backend_id) const {
     return registry_ == nullptr ? nullptr : registry_->Get(backend_id);

src/umbp/distributed/pool/peer_pool.cpp

@@ -923,6 +923,37 @@ std::vector<bool> PeerPool::BatchContains(const std::vector<std::string>& keys)
   return out;
 }
 
+std::vector<bool> PeerPool::ValidateResolved(const std::vector<uint32_t>& backend_ids,
+                                             const std::vector<std::string>& keys,
+                                             const std::vector<uint64_t>& generations) {
+  std::shared_lock<std::shared_mutex> lifecycle_lock(lifecycle_mutex_);
+  std::vector<bool> out(keys.size(), false);
+  if (keys.empty() || backends_ == nullptr) return out;
+  if (backend_ids.size() != keys.size() || generations.size() != keys.size()) return out;
+
+  // One call per backend named, not one per key.
+  std::unordered_map<uint32_t, std::vector<size_t>> by_backend;
+  for (size_t i = 0; i < keys.size(); ++i) by_backend[backend_ids[i]].push_back(i);
+
+  std::vector<std::string> batch_keys;
+  std::vector<uint64_t> batch_generations;
+  for (const auto& [backend_id, indices] : by_backend) {
+    auto* backend = backends_->Get(backend_id);
+    if (backend == nullptr) continue;  // nothing to have read from: stays false
+    batch_keys.clear();
+    batch_generations.clear();
+    for (size_t i : indices) {
+      batch_keys.push_back(keys[i]);
+      batch_generations.push_back(generations[i]);
+    }
+    const auto holds = backend->BatchStillHolds(batch_keys, batch_generations);
+    for (size_t j = 0; j < indices.size(); ++j) {
+      out[indices[j]] = j < holds.size() && holds[j];
+    }
+  }
+  return out;
+}
+
 uint32_t PeerPool::EvictSourceLocked(const PoolEvictRequest& request) const {
   constexpr uint32_t kNone = BackendRegistry::kMaxBackends;
   const auto holds = [&](uint32_t backend_id) {

src/umbp/distributed/pool_client.cpp

@@ -236,6 +236,7 @@ struct RangedPhases {
   double rmt_sub = 0.0;          // per-group sub_* vectors + PartitionBatchGetRangeTargets
   double rmt_items_build = 0.0;  // BuildContiguousToRangesItems for the scatter
   double rmt_final = 0.0;        // ApplyTransferFailures + FinalizeRemoteGetEntries
+  double rmt_check = 0.0;        // ValidateReads round trip after the RDMA read
   double rmt_medium = 0.0;       // ServeWholeObjectUnitsFromMedium
   double rmt_install = 0.0;      // install-from-arena / background prefetch
   double rmt_scatter = 0.0;      // arena -> caller buffers
@@ -311,6 +312,7 @@ struct RemoteLegSinks {
   double* peer = nullptr;
   double* decode = nullptr;
   double* final_ = nullptr;
+  double* check = nullptr;
   double* wait = nullptr;
 };
 thread_local RemoteLegSinks* t_remote_leg = nullptr;
@@ -357,6 +359,7 @@ class RangedStats {
     t.rmt_sub += p.rmt_sub;
     t.rmt_items_build += p.rmt_items_build;
     t.rmt_final += p.rmt_final;
+    t.rmt_check += p.rmt_check;
     t.rmt_medium += p.rmt_medium;
     t.rmt_install += p.rmt_install;
     t.rmt_scatter += p.rmt_scatter;
@@ -414,8 +417,8 @@ class RangedStats {
            xfer = 0, xfer_plan = 0, xfer_submit = 0, xfer_wait = 0, lock = 0, remote = 0,
            rmt_resolve = 0, rmt_build = 0, rmt_bld_items = 0, rmt_bld_plan = 0, rmt_bld_post = 0,
            rmt_peer = 0, rmt_decode = 0, rmt_sub = 0, rmt_items_build = 0, rmt_final = 0,
-           rmt_medium = 0, rmt_install = 0, rmt_wait = 0, rmt_scatter = 0, scat_plan = 0,
-           scat_submit = 0, scat_wait = 0, bytes = 0;
+           rmt_check = 0, rmt_medium = 0, rmt_install = 0, rmt_wait = 0, rmt_scatter = 0,
+           scat_plan = 0, scat_submit = 0, scat_wait = 0, bytes = 0;
   };
   struct Component {
     uint64_t calls = 0;
@@ -434,8 +437,8 @@ class RangedStats {
         "resolve={:.1f}% classify={:.1f}% build={:.1f}% validate={:.1f}% "
         "commit={:.1f}% route={:.1f}% xfer={:.1f}%(plan={:.1f}% submit={:.1f}% wait={:.1f}%) "
         "lock={:.1f}% remote={:.1f}%(rslv={:.1f}% bld={:.1f}%[it={:.1f}% pl={:.1f}% "
-        "po={:.1f}%] peer={:.1f}% dec={:.1f}% sub={:.1f}% ib={:.1f}% fin={:.1f}% med={:.1f}% "
-        "inst={:.1f}% rdma={:.1f}% "
+        "po={:.1f}%] peer={:.1f}% dec={:.1f}% sub={:.1f}% ib={:.1f}% fin={:.1f}% chk={:.1f}% "
+        "med={:.1f}% inst={:.1f}% rdma={:.1f}% "
         "scat={:.1f}%[pl={:.1f}% sub={:.1f}% wt={:.1f}%]) "
         "other={:.1f}% | xfer_only={:.2f}GiB/s end2end={:.2f}GiB/s",
         name, t.calls, t.total, 1e6 * t.total / t.calls, t.bytes / (1024.0 * 1024 * 1024),
@@ -446,9 +449,9 @@ class RangedStats {
         share(t.xfer_plan), share(t.xfer_submit), share(t.xfer_wait), share(t.lock),
         share(t.remote), share(t.rmt_resolve), share(t.rmt_build), share(t.rmt_bld_items),
         share(t.rmt_bld_plan), share(t.rmt_bld_post), share(t.rmt_peer), share(t.rmt_decode),
-        share(t.rmt_sub), share(t.rmt_items_build), share(t.rmt_final), share(t.rmt_medium),
-        share(t.rmt_install), share(t.rmt_wait), share(t.rmt_scatter), share(t.scat_plan),
-        share(t.scat_submit), share(t.scat_wait), share(other),
+        share(t.rmt_sub), share(t.rmt_items_build), share(t.rmt_final), share(t.rmt_check),
+        share(t.rmt_medium), share(t.rmt_install), share(t.rmt_wait), share(t.rmt_scatter),
+        share(t.scat_plan), share(t.scat_submit), share(t.scat_wait), share(other),
         t.xfer > 0 ? (t.bytes / t.xfer) / (1024.0 * 1024 * 1024) : 0.0,
         t.total > 0 ? (t.bytes / t.total) / (1024.0 * 1024 * 1024) : 0.0);
   }
@@ -621,6 +624,17 @@ inline std::chrono::milliseconds ResolveBusyRetryTimeout() {
   return std::chrono::milliseconds(std::min<uint64_t>(ms, 300000));
 }
 
+// Check-after-read for remote readers (#749).  On by default;
+// UMBP_REMOTE_READ_VALIDATE=0 turns it off for an A/B of its cost, which also
+// reopens the window where a read returns another key's bytes as a hit.
+inline bool RemoteReadValidateEnabled() {
+  static const bool enabled = [] {
+    const char* value = std::getenv("UMBP_REMOTE_READ_VALIDATE");
+    return value == nullptr || std::string_view(value) != "0";
+  }();
+  return enabled;
+}
+
 std::vector<PoolResolvedEntry> ResolveLocalBatchWithBusyRetry(PeerPool* pool,
                                                               const std::vector<std::string>& keys,
                                                               bool include_descs,
@@ -3864,6 +3878,7 @@ std::vector<bool> PoolClient::BatchGetRanges(const std::vector<std::string>& key
            .peer = &dbg->rmt_peer,
            .decode = &dbg->rmt_decode,
            .final_ = &dbg->rmt_final,
+           .check = &dbg->rmt_check,
            .wait = &dbg->rmt_wait};
   }
   ScopedRemoteLeg leg_scope(dbg ? &leg : nullptr);
@@ -4610,10 +4625,20 @@ void PoolClient::WaitRemoteBatchGet(RemoteGetInFlight& f, std::vector<bool>* res
     PhaseTimer wait_timer(t_remote_leg ? t_remote_leg->wait : nullptr);
     if (f.handle != nullptr) f.handle->Wait(&failures);
   }
-  PhaseTimer final_timer(t_remote_leg ? t_remote_leg->final_ : nullptr);
-  ApplyTransferFailures(f.entries, failures, "RemoteGet");
+  {
+    PhaseTimer final_timer(t_remote_leg ? t_remote_leg->final_ : nullptr);
+    ApplyTransferFailures(f.entries, failures, "RemoteGet");
+  }
+  // Before anything is reported or re-cached: a read the peer cannot vouch for
+  // must not become a hit, nor be installed locally under the key.  Timed on
+  // its own, outside `final_`, because it is a network round trip.
+  if (f.peer != nullptr) {
+    PhaseTimer check_timer(t_remote_leg ? t_remote_leg->check : nullptr);
+    ValidateRemoteGetEntries(*f.peer, f.entries);
+  }
   // Nothing to copy out here even for a staged read: the engine owns its bounce
   // pool and lands the bytes in the user's dst before its Wait returns.
+  PhaseTimer final_timer(t_remote_leg ? t_remote_leg->final_ : nullptr);
   FinalizeRemoteGetEntries(f.entries, results, recache_remote);
 }
 
@@ -4728,6 +4753,7 @@ bool PoolClient::PrepareRemoteGetEntries(const std::vector<BatchGetItem>& items,
     // once; the pages below are indices into THIS backend's buffers.
     entry.plan.backend_id = key.backend_id;
     entry.plan.pages = std::move(decoded.keys[i].pages);
+    entry.generation = key.generation;
     // Descriptors were hydrated batch-level above; the per-entry plan carries
     // none (BuildRemoteGetTransfers' EnsureBufferDescsCached call is a no-op on
     // an empty list and the read path resolves descriptors by buffer_index).
@@ -4898,6 +4924,62 @@ void PoolClient::FinalizeRemoteGetEntries(std::vector<RemoteGetEntry>& entries,
   }
 }
 
+void PoolClient::ValidateRemoteGetEntries(PeerConnection& peer,
+                                          std::vector<RemoteGetEntry>& entries) {
+  if (!RemoteReadValidateEnabled()) return;
+
+  // Ask only about reads that landed and carry a generation.  0 means the
+  // medium does not stamp copies, or the peer predates the field and so has no
+  // ValidateReads to call either.
+  ::umbp::ValidateReadsRequest req;
+  std::vector<size_t> asked;
+  for (size_t i = 0; i < entries.size(); ++i) {
+    const auto& entry = entries[i];
+    if (entry.failed || entry.generation == 0) continue;
+    req.add_keys(*entry.item->key);
+    req.add_backend_id(entry.plan.backend_id);
+    req.add_generation(entry.generation);
+    asked.push_back(i);
+  }
+  if (asked.empty()) return;
+
+  auto drop = [&](const std::vector<size_t>& which, const char* reason) {
+    for (size_t i : which) entries[i].failed = true;
+    CountMetric(MORI_UMBP_METRIC_CLIENT_REMOTE_GET_DROPPED_TOTAL,
+                MORI_UMBP_METRIC_CLIENT_REMOTE_GET_DROPPED_TOTAL_HELP, {{"reason", reason}},
+                static_cast<double>(which.size()));
+  };
+
+  // Fail closed.  A peer that stamped a generation is new enough to answer, so
+  // any error here (UNIMPLEMENTED included) means the bytes are unconfirmed,
+  // and an unconfirmed read is exactly what this check exists to refuse.
+  auto* stub = static_cast<::umbp::UMBPPeer::Stub*>(peer.peer_stub.get());
+  ::umbp::ValidateReadsResponse resp;
+  grpc::Status status(grpc::StatusCode::UNAVAILABLE, "no peer stub");
+  if (stub != nullptr) {
+    grpc::ClientContext ctx;
+    ctx.set_deadline(std::chrono::system_clock::now() + ResolveBusyRetryTimeout());
+    status = stub->ValidateReads(&ctx, req, &resp);
+  }
+  if (!status.ok() || resp.still_valid_size() != static_cast<int>(asked.size())) {
+    MORI_UMBP_WARN("[PoolClient] ValidateReads failed on {} ({} keys): {}", peer.node_id,
+                   asked.size(), status.ok() ? "answer count mismatch" : status.error_message());
+    drop(asked, "unverified");
+    return;
+  }
+
+  std::vector<size_t> stale;
+  for (size_t j = 0; j < asked.size(); ++j) {
+    if (!resp.still_valid(static_cast<int>(j))) stale.push_back(asked[j]);
+  }
+  if (stale.empty()) return;
+  MORI_UMBP_WARN(
+      "[PoolClient] BatchGet: {} of {} remote reads from {} changed under the read, returned "
+      "as misses (first key='{}')",
+      stale.size(), asked.size(), peer.node_id, *entries[stale.front()].item->key);
+  drop(stale, "stale");
+}
+
 // ---------------------------------------------------------------------------
 //  Cluster-wide existence check
 // ---------------------------------------------------------------------------

src/umbp/distributed/proto/umbp_peer.proto

@@ -24,6 +24,9 @@ service UMBPPeer {
   rpc BatchCommitSlots  (BatchCommitSlotsRequest)   returns (BatchCommitSlotsResponse);
   rpc BatchAbortSlots   (BatchAbortSlotsRequest)    returns (BatchAbortSlotsResponse);
   rpc BatchResolveKeys  (BatchResolveKeysRequest)   returns (BatchResolveKeysResponse);
+  // Check-after-read for remote readers: a reader that copied bytes out of a
+  // resolved key asks whether that copy was still the key's when it finished.
+  rpc ValidateReads     (ValidateReadsRequest)      returns (ValidateReadsResponse);
 }
 
 // ============================================================
@@ -237,4 +240,24 @@ message BatchResolveKeysResponse {
   // (buffer_index, page_index) pairs, walking keys in order.
   repeated uint32 buffer_index = 7;
   repeated uint32 page_index   = 8;
+  // Per key: which copy the pages above belong to.  A medium draws a fresh
+  // value every time it fills a slot or staging span, so the same key read
+  // twice names the same copy only if nothing was freed or reused between.
+  // The reader hands it back in ValidateReads.  0 = this medium does not stamp
+  // copies (nothing to check); absent entirely = an old peer.
+  repeated uint64 generation = 11;
+}
+
+// Parallel arrays, one entry per key read.  backend_id and generation are the
+// values the same peer returned in BatchResolveKeysResponse.
+message ValidateReadsRequest {
+  repeated string keys       = 1;
+  repeated uint32 backend_id = 2;
+  repeated uint64 generation = 3;
+}
+
+message ValidateReadsResponse {
+  // Parallel to request.keys.  False: the copy was freed or replaced while the
+  // reader was copying, so what it read may be another key's bytes.
+  repeated bool still_valid = 1;
 }