Fix(io): track RDMA notification completions in transfer status

src/io/rdma/backend_impl.cpp

@@ -815,13 +815,33 @@ NotifManager::FlushDrainStats NotifManager::ProcessOneCqe(
         struct ibv_recv_wr* bad = nullptr;
         SYSCALL_RETURN_ZERO(ibv_post_recv(ep.local.ibvHandle.qp, &wr, &bad));
       } else if (wc[i].opcode == IBV_WC_SEND) {
-        if (!IsNotifSendWrId(wc[i].wr_id)) {
-          MORI_IO_WARN(
-              "ProcessOneCqe: unexpected SEND completion with non-notification wr_id {}; "
-              "releasing 1 sqDepth under current SEND invariant",
-              wc[i].wr_id);
+        if (IsNotifSendWrId(wc[i].wr_id)) {
+          // Backward-compatible handling for any notification SENDs posted by
+          // older code that used tagged transfer IDs instead of ledger records.
+          if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed);
+        } else {
+          int mergedBatchSize = 0;
+          auto meta = ep.ledger
+                          ? ep.ledger->ReleaseByCqe(wc[i].wr_id, ep.sqDepth.get(), &mergedBatchSize)
+                          : nullptr;
+          if (meta) {
+            uint32_t finishedBefore = meta->finishedBatchSize.fetch_add(mergedBatchSize);
+            TransferStatus* statusPtr = meta->status;
+            if (statusPtr != nullptr &&
+                (finishedBefore + mergedBatchSize) == meta->totalBatchSize) {
+              statusPtr->Update(StatusCode::SUCCESS, ibv_wc_status_str(wc[i].status));
+            }
+            MORI_IO_TRACE(
+                "ProcessOneCqe: notification SEND CQE for task {} total={} finished={} cur={}",
+                meta->id, meta->totalBatchSize, finishedBefore, mergedBatchSize);
+          } else {
+            MORI_IO_WARN(
+                "ProcessOneCqe: notification SEND CQE has no ledger record for wr_id {}; "
+                "releasing 1 sqDepth under fallback path",
+                wc[i].wr_id);
+            if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed);
+          }
         }
-        if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed);
       } else {
         // Batch path: wr_id carries a recordId from the SubmissionLedger.
         uint64_t recordId = wc[i].wr_id;
@@ -1312,20 +1332,24 @@ void RdmaBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size
                                    TransferStatus* status, TransferUniqueId id, bool isRead) {
   MORI_IO_FUNCTION_TIMER;
   status->SetCode(StatusCode::IN_PROGRESS);
-  auto callbackMeta = std::make_shared<CqCallbackMeta>(status, id, 1);
+  const int notifBatchSize = config.enableNotification ? static_cast<int>(eps.size()) : 0;
+  auto callbackMeta = std::make_shared<CqCallbackMeta>(status, id, 1 + notifBatchSize);
   internal::PublishCurrentIoCallDiagnostics(callbackMeta);
 
+  RdmaTransferControl control{};
+  control.chunkBytes = config.enableTransferChunking ? config.chunkBytes : 0;
+  control.maxChunks = config.maxChunksPerTransfer;
+  control.creditByWrCount = config.enableTransferChunking;
+  control.extraCompletionCredits = notifBatchSize;
   RdmaOpRet ret = RdmaBatchReadWrite(eps, localMrPerEp, remoteMrPerEp, {localOffset},
-                                     {remoteOffset}, {size}, callbackMeta, id, isRead, 1,
-                                     config.enableTransferChunking ? config.chunkBytes : 0,
-                                     config.maxChunksPerTransfer, config.enableTransferChunking);
+                                     {remoteOffset}, {size}, callbackMeta, id, isRead, 1, control);
 
   assert(!ret.Init());
   if (ret.Failed() || ret.Succeeded()) {
     status->Update(ret.code, ret.message);
   }
   if (!ret.Failed() && config.enableNotification) {
-    RdmaOpRet notifRet = RdmaNotifyTransfer(eps, status, id);
+    RdmaOpRet notifRet = RdmaNotifyTransfer(eps, callbackMeta, id);
     if (notifRet.Failed()) {
       status->Update(notifRet.code, notifRet.message);
     }
@@ -1380,6 +1404,13 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV
     totalCredit = static_cast<int>(sizes.size());
   }
 
+  const int notifBatchSize = config.enableNotification ? static_cast<int>(eps.size()) : 0;
+  if (totalCredit > std::numeric_limits<int>::max() - notifBatchSize) {
+    status->Update(StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range");
+    return;
+  }
+  totalCredit += notifBatchSize;
+
   auto callbackMeta = std::make_shared<CqCallbackMeta>(status, id, totalCredit);
   internal::PublishCurrentIoCallDiagnostics(callbackMeta);
   RdmaOpRet ret;
@@ -1416,6 +1447,7 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV
     control.chunkBytes = chunk ? config.chunkBytes : 0;
     control.maxChunks = config.maxChunksPerTransfer;
     control.creditByWrCount = chunk;
+    control.extraCompletionCredits = notifBatchSize;
     ret = RdmaBatchReadWrite(eps, localMrPerEp, remoteMrPerEp, localOffsets, remoteOffsets, sizes,
                              callbackMeta, id, isRead, config.postBatchSize, control);
   }
@@ -1424,7 +1456,7 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV
     status->Update(ret.code, ret.message);
   }
   if (!ret.Failed() && config.enableNotification) {
-    RdmaOpRet notifRet = RdmaNotifyTransfer(eps, status, id);
+    RdmaOpRet notifRet = RdmaNotifyTransfer(eps, callbackMeta, id);
     if (notifRet.Failed()) {
       status->Update(notifRet.code, notifRet.message);
     }

src/io/rdma/common.cpp

@@ -395,9 +395,9 @@ static void ResetMergedWorkRequestPointers(MergedWorkRequest* wr) {
 /*                                         Rdma Utilities                                         */
 /* ---------------------------------------------------------------------------------------------- */
 
-RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, TransferUniqueId id) {
+RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, std::shared_ptr<CqCallbackMeta> callbackMeta,
+                             TransferUniqueId id) {
   MORI_IO_FUNCTION_TIMER;
-  (void)status;
 
   std::string reserveErr;
   int reserved = 0;
@@ -413,6 +413,12 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans
   }
 
   for (size_t i = 0; i < eps.size(); i++) {
+    if (!eps[i].ledger) {
+      for (int j = static_cast<int>(i); j < reserved; ++j) ReleaseSqDepth(eps[j], 1);
+      return {StatusCode::ERR_RDMA_OP,
+              "submission ledger is not initialized for notification SEND tracking"};
+    }
+
     const application::RdmaEndpoint& ep = eps[i].local;
     NotifMessage msg{id, static_cast<int>(i), static_cast<int>(eps.size())};
 
@@ -421,8 +427,10 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans
     sge.length = sizeof(NotifMessage);
     sge.lkey = 0;
 
+    const uint64_t recordId = eps[i].ledger->Insert(1, true, callbackMeta, 1);
+
     struct ibv_send_wr wr{};
-    wr.wr_id = MakeNotifSendWrId(id);
+    wr.wr_id = recordId;
     wr.opcode = IBV_WR_SEND;
     wr.send_flags = IBV_SEND_INLINE | IBV_SEND_SIGNALED;
     wr.sg_list = &sge;
@@ -431,8 +439,11 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans
     struct ibv_send_wr* bad_wr = nullptr;
     int ret = ibv_post_send(ep.ibvHandle.qp, &wr, &bad_wr);
     if (ret != 0) {
-      // WR i was reserved but failed to post if bad_wr points at this WR.
-      if (bad_wr == &wr) ReleaseSqDepth(eps[i], 1);
+      // This call posts a single WR, so a non-zero return means the current
+      // notification was not accepted and no CQE will arrive for its ledger record.
+      ReleaseSqDepth(eps[i], 1);
+      int dummy = 0;
+      eps[i].ledger->ReleaseByCqe(recordId, nullptr, &dummy);
       // Any remaining endpoints are reserved but not posted yet.
       for (int j = i + 1; j < eps.size(); ++j) ReleaseSqDepth(eps[j], 1);
       std::string message =
@@ -705,7 +716,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps,
       return {StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range"};
     }
     for (size_t k = 0; k < mergedWrCount; ++k) mergedWrs[k].mergedRequests = 1;
-    if (control.ownsTotalBatchSize) callbackMeta->totalBatchSize = static_cast<int>(mergedWrCount);
+    if (control.ownsTotalBatchSize) {
+      callbackMeta->totalBatchSize =
+          static_cast<int>(mergedWrCount) + control.extraCompletionCredits;
+    }
   }
 
   size_t epNum = eps.size();

src/io/rdma/common.hpp

@@ -253,14 +253,16 @@ struct RdmaOpRet {
   bool Failed() { return code > StatusCode::ERR_BEGIN; }
 };
 
-RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, TransferUniqueId id);
+RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, std::shared_ptr<CqCallbackMeta> callbackMeta,
+                             TransferUniqueId id);
 
 struct RdmaTransferControl {
   size_t chunkBytes{0};
   int maxChunks{1};
   bool creditByWrCount{false};
   bool ownsTotalBatchSize{true};
   bool disableMerge{false};
+  int extraCompletionCredits{0};
 };
 
 RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps,

src/umbp/distributed/master/global_block_index.cpp

@@ -130,6 +130,9 @@ size_t ApplyAddOrRemoveLocked(EntryMap& entries, NodeToKeys& node_to_keys,
       return 0;
     }
     loc->size = ev.size;
+    loc->encoding = ev.encoding;
+    if (loc->encoding.stored_bytes == 0) loc->encoding.stored_bytes = ev.size;
+    if (loc->encoding.logical_bytes == 0) loc->encoding.logical_bytes = loc->encoding.stored_bytes;
     return 1;
   }
   // REMOVE
@@ -222,6 +225,9 @@ void GlobalBlockIndex::ReplaceNodeLocations(const std::string& node_id,
       auto [loc, inserted] = FindOrInsertLocation(entry, node_id, ev->tier);
       (void)inserted;
       loc->size = ev->size;
+      loc->encoding = ev->encoding;
+      if (loc->encoding.stored_bytes == 0) loc->encoding.stored_bytes = ev->size;
+      if (loc->encoding.logical_bytes == 0) loc->encoding.logical_bytes = loc->encoding.stored_bytes;
       sh.node_to_keys[node_id].insert(ev->key);
     }
   }

src/umbp/distributed/master/master_client.cpp

@@ -35,6 +35,7 @@
 #include "umbp/common/env_time.h"
 #include "umbp/distributed/master/master_metrics.h"
 #include "umbp/distributed/master/rpc_latency_timer.h"
+#include "umbp/distributed/peer/batch_resolve_codec.h"
 #include "umbp/distributed/peer/peer_dram_allocator.h"
 #include "umbp/distributed/peer/peer_ssd_manager.h"
 
@@ -108,6 +109,7 @@ void FillBundle(::umbp::EventBundle* dst, const EventBundle& src) {
     pe->set_key(ev.key);
     pe->set_tier(ToProtoTier(ev.tier));
     pe->set_size(ev.size);
+    FillProtoKvEncoding(ev.encoding, pe->mutable_encoding());
   }
 }
 
@@ -258,6 +260,7 @@ grpc::Status MasterClient::RouteGet(const std::string& key,
   r.node_id = resp.node_id();
   r.tier = FromProtoTier(resp.tier());
   r.size = resp.size();
+  r.encoding = KvEncodingFromProto(resp.encoding(), r.size);
   r.peer_address = resp.peer_address();
   *out_result = std::move(r);
   return grpc::Status::OK;
@@ -333,7 +336,8 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector<std::string>& keys,
   // Columnar response: node_ref[i] is a 1-based index into resp.nodes()
   // (0 = not found); tier[i] / size[i] are parallel per-key arrays.
   const int n = resp.node_ref_size();
-  if (resp.tier_size() != n || resp.size_size() != n) {
+  if (resp.tier_size() != n || resp.size_size() != n ||
+      (resp.encoding_size() != 0 && resp.encoding_size() != n)) {
     return grpc::Status(grpc::StatusCode::INTERNAL,
                         "BatchRouteGet: malformed columnar response (array length mismatch)");
   }
@@ -349,6 +353,8 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector<std::string>& keys,
     r.node_id = node.node_id();
     r.tier = FromProtoTier(resp.tier(i));
     r.size = resp.size(i);
+    r.encoding =
+        resp.encoding_size() == n ? KvEncodingFromProto(resp.encoding(i), r.size) : RawKvEncoding(r.size);
     r.peer_address = node.peer_address();
     (*out)[i] = std::move(r);
   }

src/umbp/distributed/master/master_server.cpp

@@ -32,6 +32,7 @@
 #include <vector>
 
 #include "mori/utils/mori_log.hpp"
+#include "umbp/distributed/peer/batch_resolve_codec.h"
 #include "umbp.grpc.pb.h"
 #include "umbp/common/env_time.h"
 #include "umbp/distributed/master/external_kv_block_index.h"
@@ -98,6 +99,7 @@ KvEvent FromProtoEvent(const ::umbp::KvEvent& pe) {
   ev.key = pe.key();
   ev.tier = static_cast<TierType>(pe.tier());
   ev.size = pe.size();
+  ev.encoding = KvEncodingFromProto(pe.encoding(), ev.size);
   return ev;
 }
 
@@ -376,6 +378,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser
     response->set_tier(static_cast<::umbp::TierType>(result->location.tier));
     response->set_size(result->location.size);
     response->set_peer_address(result->peer_address);
+    FillProtoKvEncoding(result->location.encoding, response->mutable_encoding());
 
     if (metrics_) {
       metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_GET,
@@ -439,6 +442,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser
         response->add_node_ref(0);
         response->add_tier(::umbp::TIER_UNKNOWN);
         response->add_size(0);
+        FillProtoKvEncoding(RawKvEncoding(0), response->add_encoding());
         continue;
       }
       std::string node_key = opt->location.node_id;
@@ -454,6 +458,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser
       response->add_node_ref(it->second);
       response->add_tier(static_cast<::umbp::TierType>(opt->location.tier));
       response->add_size(opt->location.size);
+      FillProtoKvEncoding(opt->location.encoding, response->add_encoding());
       if (metrics_) {
         metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET,
                              MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET_HELP,

src/umbp/distributed/peer/peer_dram_allocator.cpp

@@ -28,6 +28,7 @@
 #include <utility>
 
 #include "mori/utils/mori_log.hpp"
+#include "umbp/codec/kv_encoding.h"
 
 namespace mori::umbp {
 
@@ -101,12 +102,18 @@ const PageBitmapAllocator* PeerDramAllocator::AllocatorForLocked(TierType tier)
 
 PeerDramAllocator::AllocateResult PeerDramAllocator::Allocate(const std::string& key, uint64_t size,
                                                               TierType tier) {
+  return Allocate(key, size, tier, RawKvEncoding(size));
+}
+
+PeerDramAllocator::AllocateResult PeerDramAllocator::Allocate(
+    const std::string& key, uint64_t size, TierType tier, const KvEncodingDescriptor& encoding) {
   std::lock_guard<std::mutex> lock(mutex_);
-  return AllocateLocked(key, size, tier);
+  return AllocateLocked(key, size, tier, encoding);
 }
 
 PeerDramAllocator::AllocateResult PeerDramAllocator::AllocateLocked(const std::string& key,
-                                                                    uint64_t size, TierType tier) {
+                                                                    uint64_t size, TierType tier,
+                                                                    const KvEncodingDescriptor& encoding) {
   auto fail = [&](Outcome outcome, const char* reason) {
     MORI_UMBP_WARN("[PeerDramAllocator] Allocate reason={} key='{}' size={} tier={}", reason, key,
                    size, static_cast<int>(tier));
@@ -145,6 +152,9 @@ PeerDramAllocator::AllocateResult PeerDramAllocator::AllocateLocked(const std::s
   slot.tier = tier;
   slot.pages = std::move(*pages);
   slot.size = size;
+  slot.encoding = encoding;
+  if (slot.encoding.stored_bytes == 0) slot.encoding.stored_bytes = size;
+  if (slot.encoding.logical_bytes == 0) slot.encoding.logical_bytes = slot.encoding.stored_bytes;
   slot.deadline = std::chrono::steady_clock::now() + pending_ttl_;
   slot.generation = allocator_generation_;
   pending_[slot.slot_id] = slot;
@@ -160,7 +170,7 @@ std::vector<PeerDramAllocator::BatchAllocateResult> PeerDramAllocator::BatchAllo
   std::lock_guard<std::mutex> lock(mutex_);
   for (size_t i = 0; i < entries.size(); ++i) {
     const auto& entry = entries[i];
-    auto result = AllocateLocked(entry.key, entry.size, entry.tier);
+    auto result = AllocateLocked(entry.key, entry.size, entry.tier, entry.encoding);
     out[i].outcome = result.outcome;
     out[i].slot = std::move(result.slot);
     if (out[i].outcome == Outcome::kSuccessAllocated && out[i].slot.has_value()) {
@@ -219,7 +229,8 @@ bool PeerDramAllocator::CommitLocked(uint64_t slot_id, const std::string& key,
   owned.tier = it->second.tier;
   owned.pages = std::move(it->second.pages);
   owned.size = it->second.size;
-  QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, owned.tier, owned.size});
+  owned.encoding = it->second.encoding;
+  QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, owned.tier, owned.size, owned.encoding});
   owned_[key] = std::move(owned);
   pending_.erase(it);
   bytes_committed = owned_[key].size;
@@ -276,6 +287,7 @@ PeerDramAllocator::ResolveResult PeerDramAllocator::Resolve(const std::string& k
   r.tier = it->second.tier;
   r.pages = it->second.pages;
   r.size = it->second.size;
+  r.encoding = it->second.encoding;
   // Extend the read lease so concurrent Evict reports bytes_freed=0 for
   // this key.  steady_clock is monotonic and read_lease_ttl_ is fixed,
   // so this assignment is always >= any previous deadline for the key.
@@ -297,6 +309,7 @@ std::vector<PeerDramAllocator::ResolvedEntry> PeerDramAllocator::BatchResolve(
     entry.tier = it->second.tier;
     entry.pages = it->second.pages;
     entry.size = it->second.size;
+    entry.encoding = it->second.encoding;
     if (include_descs) {
       entry.descs = BuildBufferDescsLocked(it->second.tier, it->second.pages);
     }
@@ -360,6 +373,7 @@ std::optional<PeerDramAllocator::DramCopyPin> PeerDramAllocator::AcquireDramCopy
 
   DramCopyPin pin;
   pin.total_size = it->second.size;
+  pin.encoding = it->second.encoding;
   pin.segments = BuildCopySegmentsLocked(it->second.tier, it->second.pages, it->second.size);
   pin.pin_token = next_pin_token_++;
   pins_[key] = PinState{pin.pin_token, std::chrono::steady_clock::now()};
@@ -496,7 +510,8 @@ std::vector<KvEvent> PeerDramAllocator::SnapshotOwnedKeysLocked() const {
   std::vector<KvEvent> out;
   out.reserve(owned_.size());
   for (const auto& kv : owned_) {
-    out.push_back(KvEvent{KvEvent::Kind::ADD, kv.first, kv.second.tier, kv.second.size});
+    out.push_back(
+        KvEvent{KvEvent::Kind::ADD, kv.first, kv.second.tier, kv.second.size, kv.second.encoding});
   }
   return out;
 }