perf(umbp): send and resolve a key set once, not once per layer group

src/umbp/distributed/pool_client.cpp

@@ -1488,15 +1488,24 @@ void PoolClient::ResolveLocalBatch(const std::vector<std::string>& keys,
   if (resolutions != nullptr) resolutions->assign(keys.size(), ResolvedEntry{});
   if (candidates.empty() || default_pool_ == nullptr || registry_.Empty()) return;
 
-  std::vector<std::string> batch;
-  batch.reserve(candidates.size());
-  for (size_t i : candidates) batch.push_back(keys[i]);
+  // A whole batch -- every key a candidate, in order -- is what a restore asks
+  // for, and there `keys` already IS the query.  Naming it directly saves
+  // copying the entire key set, which for a layer-wise reader is a thousand-odd
+  // ~128-byte strings rebuilt once per layer group for a query it just made.
+  // A partial batch still has to be gathered.
+  std::vector<std::string> gathered;
+  const std::vector<std::string>* batch = &keys;
+  if (candidates.size() != keys.size()) {
+    gathered.reserve(candidates.size());
+    for (size_t i : candidates) gathered.push_back(keys[i]);
+    batch = &gathered;
+  }
 
   // Both readers fed from here — ServeLocalGets and BatchGetRanges — can
   // consume a file ref, and this resolve never leaves the process, so ask for
   // one whenever the caller can say where the bytes land.
   const bool want_file_refs = static_cast<bool>(dst_is_device);
-  auto found = ResolveLocalBatchWithBusyRetry(default_pool_.get(), batch,
+  auto found = ResolveLocalBatchWithBusyRetry(default_pool_.get(), *batch,
                                               /*include_descs=*/false,
                                               /*allow_file_refs=*/want_file_refs);
 
@@ -1513,7 +1522,7 @@ void PoolClient::ResolveLocalBatch(const std::vector<std::string>& keys,
     for (size_t j = 0; j < candidates.size() && j < found.size(); ++j) {
       if (!found[j].resolved.file_ref.IsFile()) continue;
       if (dst_is_device(candidates[j])) continue;
-      restage.push_back(batch[j]);
+      restage.push_back((*batch)[j]);
       restage_at.push_back(j);
     }
     if (!restage.empty()) {

src/umbp/distributed/proto/umbp_standalone.proto

@@ -101,6 +101,19 @@ message BatchRangeDataRequest {
   repeated uint64 object_offsets = 7;
   // Put only: the expected full object size, one per key. Empty for get.
   repeated uint64 object_sizes = 8;
+  // Get only, optional. A layer-wise reader names one key set once per layer
+  // group, changing only which bytes it wants, so the keys are the one part of
+  // the request that does not change between the calls. Non-zero key_handle
+  // names a key list the server minted on an earlier call and still holds, and
+  // `keys` is then empty.
+  //
+  // This is a hint and nothing more: a handle the server no longer holds --
+  // evicted, or from an earlier server -- is answered with key_handle_unknown
+  // and the caller simply repeats the call carrying its keys. key_fingerprint
+  // is checked against the one recorded when the handle was minted, so a handle
+  // can never bind to a key list other than the one it was minted for.
+  uint64 key_handle = 9;
+  uint64 key_fingerprint = 10;
 }
 
 message BatchDataWithDepthRequest {
@@ -114,6 +127,14 @@ message BatchDataWithDepthRequest {
 
 message BatchBoolResponse {
   repeated bool ok = 1;
+  // BatchGetRanges only. Non-zero when the server took this call's key list and
+  // minted a handle for it; the caller may name that handle instead of the keys
+  // next time it asks about the same set.
+  uint64 key_handle = 2;
+  // The offered key_handle is not one the server holds, or does not match the
+  // offered fingerprint. Nothing was read; `ok` is empty. Repeat the call with
+  // the keys.
+  bool key_handle_unknown = 3;
 }
 
 enum MemoryKind {

src/umbp/include/umbp/standalone/standalone_process_client.h

@@ -29,6 +29,7 @@
 #include <cstdint>
 #include <memory>
 #include <mutex>
+#include <random>
 #include <shared_mutex>
 #include <string>
 #include <vector>
@@ -138,6 +139,58 @@ class StandaloneProcessClient : public IUMBPClient {
   mutable std::mutex registration_mu_;
   std::string client_id_;
   std::vector<RegisteredRegion> regions_;
+
+  // The key lists this client has already sent, and the handles the server
+  // gave back for them.
+  //
+  // A layer-wise restore asks about one key set once per layer group, so the
+  // keys are the only part of a ranged get that does not change between those
+  // calls -- and at a thousand-odd ~128-byte keys they are the expensive part
+  // to serialize. Naming a remembered list instead costs a comparison against
+  // what was sent last time.
+  //
+  // Matching is by full equality of the key vector, not by hash, so nothing
+  // here can select the wrong list: the fingerprint travels only so the server
+  // can make the same check on its side. A handle the server has dropped comes
+  // back as key_handle_unknown, the entry is discarded, and the call is simply
+  // repeated carrying the keys.
+  struct KeyHandle {
+    std::vector<std::string> keys;
+    uint64_t handle = 0;
+    uint64_t fingerprint = 0;
+  };
+
+  // A restore does NOT have one key set in flight per pool, which is what this
+  // held eight slots for. A layer-wise reader splits a pool's keys to fit a
+  // per-call range budget and walks the pieces in order, once per layer group,
+  // so the sets arrive as a fixed-order cycle. That is the one access pattern
+  // LRU cannot serve at all: below the cycle length the entry it evicts is
+  // always the next one wanted, so the hit rate is 0 rather than degraded.
+  // Measured on the wire bench at 512-key chunks: 100% at eight sets, 0% at
+  // nine, with 495 of 504 mints re-minting a set that had just been used.
+  //
+  // Replacement is therefore random. No access order can drive it into that
+  // failure, and past capacity it decays as capacity/cycle instead of falling
+  // off a cliff -- which matters because the cycle length belongs to the
+  // caller's chunking, and is not something this side gets to know.
+  //
+  // Capacity is a guess at that chunking, so UMBP_KEY_HANDLE_SLOTS overrides
+  // it. Each slot holds one key set: for the thousand-odd 128-byte keys a
+  // layer group asks about, roughly 160 KiB.
+  static size_t KeyHandleSlots();
+
+  // Returns 0 when this set has not been sent before, and fills *fingerprint
+  // either way.
+  uint64_t LookupKeyHandle(const std::vector<std::string>& keys, uint64_t* fingerprint);
+  void RememberKeyHandle(const std::vector<std::string>& keys, uint64_t handle,
+                         uint64_t fingerprint);
+  void ForgetKeyHandle(uint64_t handle);
+
+  std::mutex key_handle_mu_;
+  std::vector<KeyHandle> key_handles_;  // unordered: the victim is drawn, not aged
+  // Only has to be uncorrelated with the caller's access order, so a fixed
+  // seed is deliberate: it keeps a run reproducible.
+  std::minstd_rand key_handle_rng_{0x9e3779b9};
 };
 
 }  // namespace mori::umbp::standalone

src/umbp/standalone/standalone_process_client.cpp

@@ -52,6 +52,31 @@ namespace {
 
 std::atomic<uint64_t> g_client_counter{0};
 
+// Travels with a ranged get that names its keys by handle, so the server can
+// check the handle still stands for the list the caller means. It is a check,
+// not the lookup: both sides also hold the keys themselves and the client side
+// compares them in full, so a collision costs a resend and never a wrong read.
+//
+// Computed once per key set, on the call that introduces it -- the same call
+// that pays to serialize the keys anyway -- and reused for every call after.
+// Zero is reserved to mean "do not bother remembering this set".
+uint64_t FingerprintKeys(const std::vector<std::string>& keys) {
+  constexpr uint64_t kOffsetBasis = 1469598103934665603ULL;
+  constexpr uint64_t kPrime = 1099511628211ULL;
+  uint64_t hash = kOffsetBasis;
+  const auto mix = [&hash](unsigned char byte) {
+    hash ^= byte;
+    hash *= kPrime;
+  };
+  for (const std::string& key : keys) {
+    // Length-delimited, so that concatenations that happen to agree do not.
+    const uint64_t length = key.size();
+    for (int shift = 0; shift < 64; shift += 8) mix(static_cast<unsigned char>(length >> shift));
+    for (char c : key) mix(static_cast<unsigned char>(c));
+  }
+  return hash == 0 ? kPrime : hash;
+}
+
 ::umbp::TierType TierToProto(TierType tier) {
   switch (tier) {
     case TierType::HBM:
@@ -430,6 +455,81 @@ std::vector<bool> StandaloneProcessClient::BatchGet(const std::vector<std::strin
   return std::vector<bool>(resp.ok().begin(), resp.ok().end());
 }
 
+uint64_t StandaloneProcessClient::LookupKeyHandle(const std::vector<std::string>& keys,
+                                                  uint64_t* fingerprint) {
+  // Nothing to name, and nothing worth remembering: a zero fingerprint tells
+  // the server not to mint a handle for it.
+  if (keys.empty()) {
+    *fingerprint = 0;
+    return 0;
+  }
+  // Same signal when this client keeps no handles at all. Without it the switch
+  // is only half a switch: RememberKeyHandle drops the handle, but the server
+  // has already been told to keep the key set alive for a client that will
+  // never ask for it -- a table of them, at the server's whole capacity. It
+  // also skips fingerprinting every key on a call that has no use for the
+  // answer.
+  if (KeyHandleSlots() == 0) {
+    *fingerprint = 0;
+    return 0;
+  }
+  {
+    std::lock_guard<std::mutex> lock(key_handle_mu_);
+    for (size_t i = 0; i < key_handles_.size(); ++i) {
+      // Size, then the two ends, then the whole thing. The cheap checks reject
+      // a different set outright; the full compare runs only for the set that
+      // is about to be a hit, and is what makes the match exact.
+      const KeyHandle& entry = key_handles_[i];
+      if (entry.keys.size() != keys.size()) continue;
+      if (entry.keys.front() != keys.front() || entry.keys.back() != keys.back()) continue;
+      if (entry.keys != keys) continue;
+      *fingerprint = entry.fingerprint;
+      return entry.handle;
+    }
+  }
+  *fingerprint = FingerprintKeys(keys);
+  return 0;
+}
+
+size_t StandaloneProcessClient::KeyHandleSlots() {
+  static const size_t slots = [] {
+    // Enough for a layer group's worth of chunks several times over. The cost
+    // of being too small is a resend, the cost of being too large is memory,
+    // and only one of those is recoverable at runtime.
+    size_t configured = 128;
+    if (const char* raw = std::getenv("UMBP_KEY_HANDLE_SLOTS")) {
+      char* end = nullptr;
+      const unsigned long long parsed = std::strtoull(raw, &end, 10);
+      if (end != raw && *end == '\0') configured = static_cast<size_t>(parsed);
+    }
+    return configured;
+  }();
+  return slots;
+}
+
+void StandaloneProcessClient::RememberKeyHandle(const std::vector<std::string>& keys,
+                                                uint64_t handle, uint64_t fingerprint) {
+  if (handle == 0) return;
+  const size_t slots = KeyHandleSlots();
+  if (slots == 0) return;  // a way to turn the mechanism off outright
+  std::lock_guard<std::mutex> lock(key_handle_mu_);
+  if (key_handles_.size() < slots) {
+    key_handles_.push_back(KeyHandle{keys, handle, fingerprint});
+    return;
+  }
+  // Full: draw the victim rather than dropping the oldest. See the header --
+  // the sets arrive as a fixed-order cycle, and evicting by age under a cycle
+  // means evicting exactly the set about to be asked for.
+  key_handles_[key_handle_rng_() % key_handles_.size()] = KeyHandle{keys, handle, fingerprint};
+}
+
+void StandaloneProcessClient::ForgetKeyHandle(uint64_t handle) {
+  std::lock_guard<std::mutex> lock(key_handle_mu_);
+  key_handles_.erase(std::remove_if(key_handles_.begin(), key_handles_.end(),
+                                    [handle](const KeyHandle& e) { return e.handle == handle; }),
+                     key_handles_.end());
+}
+
 std::vector<bool> StandaloneProcessClient::BatchGetRanges(
     const std::vector<std::string>& keys, const std::vector<std::vector<uintptr_t>>& dsts,
     const std::vector<std::vector<size_t>>& sizes,
@@ -439,13 +539,19 @@ std::vector<bool> StandaloneProcessClient::BatchGetRanges(
   std::shared_lock lk(op_mutex_);
   if (closed_ || !RangeBatchShapeValid(keys.size(), dsts, sizes, src_offsets)) return failed;
 
+  uint64_t fingerprint = 0;
+  uint64_t handle = LookupKeyHandle(keys, &fingerprint);
+
   ::umbp::BatchRangeDataRequest req;
   req.set_client_id(ClientId());
+  req.set_key_fingerprint(fingerprint);
+  // The keys are deliberately NOT added here: whether they travel at all is
+  // decided per attempt below, and registration_mu_ guards only the region
+  // lookup that OffsetForLocked does.
   {
     std::lock_guard<std::mutex> registration_lock(registration_mu_);
     for (size_t i = 0; i < keys.size(); ++i) {
       if (dsts[i].size() > std::numeric_limits<uint32_t>::max()) return failed;
-      req.add_keys(keys[i]);
       req.add_range_counts(static_cast<uint32_t>(dsts[i].size()));
       for (size_t j = 0; j < dsts[i].size(); ++j) {
         uint64_t shm_offset = 0;
@@ -459,11 +565,33 @@ std::vector<bool> StandaloneProcessClient::BatchGetRanges(
     }
   }
 
-  grpc::ClientContext ctx;
-  ::umbp::BatchBoolResponse resp;
-  const grpc::Status status = stub_->BatchGetRanges(&ctx, req, &resp);
-  if (!status.ok() || resp.ok_size() != static_cast<int>(keys.size())) return failed;
-  return std::vector<bool>(resp.ok().begin(), resp.ok().end());
+  // At most twice: once naming a handle, and if the server no longer holds it,
+  // once carrying the keys. A handle is only ever offered after the server
+  // handed it out, so the retry is the rare path.
+  for (int attempt = 0; attempt < 2; ++attempt) {
+    if (handle != 0) {
+      req.set_key_handle(handle);
+      req.clear_keys();
+    } else {
+      req.set_key_handle(0);
+      req.mutable_keys()->Reserve(static_cast<int>(keys.size()));
+      for (const auto& key : keys) req.add_keys(key);
+    }
+
+    grpc::ClientContext ctx;
+    ::umbp::BatchBoolResponse resp;
+    const grpc::Status status = stub_->BatchGetRanges(&ctx, req, &resp);
+    if (!status.ok()) return failed;
+    if (resp.key_handle_unknown()) {
+      ForgetKeyHandle(handle);
+      handle = 0;
+      continue;
+    }
+    if (resp.ok_size() != static_cast<int>(keys.size())) return failed;
+    RememberKeyHandle(keys, resp.key_handle(), fingerprint);
+    return std::vector<bool>(resp.ok().begin(), resp.ok().end());
+  }
+  return failed;
 }
 
 std::vector<bool> StandaloneProcessClient::BatchPutRanges(