[peregrine] Add ResolveAddrCallback and base-address resolution in RequestTracker.

PiperOrigin-RevId: 1008474233

Author: mubashirq

peregrine/src/api/transport_test.cc

@@ -58,17 +58,23 @@ class TransportTest : public ::testing::TestWithParam<Param> {
         completed_status_(Status::kNotFound),
         partial_(kBufSize / 2),
         l_(kBufSize, kNumConnsPerPeer),
-        r_(kBufSize, kNumConnsPerPeer,
-           [this](Status s, const Metadata& post_meta,
-                  absl::Span<const Metadata> req_metas) {
-             absl::MutexLock lock(cb_mu_);
-             recv_posts_.push_back(PostEvent{
-                 .status = s,
-                 .post_meta = post_meta,
-                 .req_metas =
-                     std::vector<Metadata>(req_metas.begin(), req_metas.end()),
-             });
-           }) {
+        r_(
+            kBufSize, kNumConnsPerPeer,
+            [this](Status s, const Metadata& post_meta,
+                   absl::Span<const Metadata> req_metas) {
+              absl::MutexLock lock(cb_mu_);
+              recv_posts_.push_back(PostEvent{
+                  .status = s,
+                  .post_meta = post_meta,
+                  .req_metas =
+                      std::vector<Metadata>(req_metas.begin(), req_metas.end()),
+              });
+            },
+            [this](const Metadata& req_meta) -> Byte* {
+              if (req_meta == Metadata{1, 2, 3}) return r_.DataPtr();
+              if (req_meta == Metadata{4, 5, 6}) return r_.DataPtr() + partial_;
+              return nullptr;
+            }) {
     CHECK(1 <= partial_ && partial_ < l_.DataSize());
     DCHECK_EQ(l_.DataSize(), r_.DataSize());
   }
@@ -172,7 +178,7 @@ TEST_P(TransportTest, Write) {
   const Request req2 = {
       .op = Op::kWrite,
       .laddr = l_.DataPtr() + partial_,
-      .raddr = r_.DataPtr() + partial_,
+      .raddr = nullptr,
       .len = l_.DataSize() - partial_,
       .metadata = {4, 5, 6},
   };

peregrine/src/api/transport_types.h

@@ -51,9 +51,7 @@ struct Request final {
   Metadata metadata = {};  // optional opaque per-request metadata
 
   // Returns true iff the request is valid.
-  constexpr bool IsValid() const {
-    return laddr != nullptr && raddr != nullptr && 1 <= len;
-  }
+  constexpr bool IsValid() const { return laddr != nullptr && 1 <= len; }
 
   // Returns true iff the requests are equal.
   friend bool operator==(const Request& a, const Request& b) {
@@ -85,6 +83,11 @@ using RecvPostCallback =
     absl::AnyInvocable<void(Status, const Metadata& post_metadata,
                             absl::Span<const Metadata> req_metadata) const>;
 
+// Callback invoked on the receiver to resolve the target buffer address for an
+// incoming request when `Request::raddr` is not provided by the sender.
+using ResolveAddrCallback =
+    absl::AnyInvocable<Byte*(const Metadata& req_metadata) const>;
+
 // Returns true iff the transport operation is still in progress.
 constexpr bool IsInProgress(const Status s) { return s == Status::kInProgress; }
 

peregrine/src/api/transport_types_test.cc

@@ -62,6 +62,15 @@ TEST(TransportRequest, Validity) {
   EXPECT_TRUE(rv.IsValid());
   EXPECT_TRUE(IsValid({rv}));
   LOG(INFO) << rv;
+
+  const Request rv_null_raddr = {
+      .op = Op::kWrite,
+      .laddr = reinterpret_cast<Byte*>(0x1234),
+      .raddr = nullptr,
+      .len = 1,
+  };
+  EXPECT_TRUE(rv_null_raddr.IsValid());
+  EXPECT_TRUE(IsValid({rv_null_raddr}));
 }
 
 TEST(TransportRequest, RKeyAndEquality) {

peregrine/src/api/transport_util.cc

@@ -37,7 +37,8 @@ std::unique_ptr<Transport> CreateTransport(std::string_view endpoint,
   };
 
   return TransportImpl::Create(config, e, std::move(creds),
-                               std::move(options.recv_post_callback));
+                               std::move(options.recv_post_callback),
+                               std::move(options.resolve_addr_callback));
 }
 
 }  // namespace peregrine

peregrine/src/api/transport_util.h

@@ -20,6 +20,11 @@ struct Options {
   // Callback invoked on the receiver when an incoming `Post()` batch completes.
   // If `nullptr`, no callback is invoked on post receive completion.
   RecvPostCallback recv_post_callback = nullptr;
+
+  // Callback invoked on the receiver to resolve the target buffer address for
+  // an incoming request when `Request::raddr` is not provided by the sender.
+  // If `nullptr`, incoming requests with a null `raddr` are rejected.
+  ResolveAddrCallback resolve_addr_callback = nullptr;
 };
 
 // Creates a new transport instance that listens for control messages on

peregrine/src/internal/chunk/header.h

@@ -122,9 +122,6 @@ struct alignas(8) DataHeader final : public HeaderBase {
   // Returns true iff this `DataHeader` is valid.
   bool IsValid() const { return IsData() || IsAck(); }
 
-  // Returns the destination memory address for this chunk.
-  Byte* DstAddr() const { return reinterpret_cast<Byte*>(raddr.value()); }
-
   // Returns a string representation for debugging/logging.
   std::string ToString() const;
 
@@ -189,8 +186,8 @@ struct alignas(8) RequestInfo final {
 
   // Returns true iff this `RequestInfo` is valid.
   bool IsValid() const {
-    return (op == Op::kRead || op == Op::kWrite) && raddr.value() != 0 &&
-           total_len > 0 && nchunks > 0 && recipe.IsValid();
+    return (op == Op::kRead || op == Op::kWrite) && total_len > 0 &&
+           nchunks > 0 && recipe.IsValid();
   }
 
   // Returns a string representation for debugging/logging.

peregrine/src/internal/chunk/header_test.cc

@@ -42,7 +42,6 @@ TEST(HeaderTest, PostInfoDataAndAckProperties) {
   EXPECT_TRUE(data.IsData());
   EXPECT_FALSE(data.IsAck());
   EXPECT_TRUE(data.IsValid());
-  EXPECT_EQ(data.DstAddr(), reinterpret_cast<Byte*>(kBufferBaseAddr.value()));
   EXPECT_TRUE(IsMatch(data, GenPayload(kChunkSize)));
   EXPECT_FALSE(IsMatch(data, GenPayload(kChunkSize - 1)));
   EXPECT_EQ(data, data);
@@ -75,6 +74,9 @@ TEST(RequestInfoTest, ValidationAndEquality) {
   EXPECT_TRUE(req.IsValid());
   EXPECT_EQ(req, req);
   LOG(INFO) << req;
+
+  req.raddr = addr_t(0);
+  EXPECT_TRUE(req.IsValid());
 }
 
 }  // namespace

peregrine/src/internal/engine/engine.cc

@@ -51,9 +51,10 @@ absl::Status NotFoundError(const Handle h) {
 }
 }  // namespace
 
-std::unique_ptr<Engine> Engine::Create(const Config& config, HostInfo& self,
-                                       Control& control,
-                                       RecvPostCallback recv_post_callback) {
+std::unique_ptr<Engine> Engine::Create(
+    const Config& config, HostInfo& self, Control& control,
+    RecvPostCallback recv_post_callback,
+    ResolveAddrCallback resolve_addr_callback) {
   static_assert(assumptions::kHostInfoDependsOnControlAndDataPlanes);
 
   if ABSL_PREDICT_FALSE (!config.IsValid()) {
@@ -67,20 +68,23 @@ std::unique_ptr<Engine> Engine::Create(const Config& config, HostInfo& self,
   }
   DCHECK(self.IsDataPlaneValid());
   return absl::WrapUnique(new Engine(config, self, std::move(helper),
-                                     std::move(recv_post_callback)));
+                                     std::move(recv_post_callback),
+                                     std::move(resolve_addr_callback)));
 }
 
 Engine::Engine(const Config& config, const HostInfo& self,
                std::unique_ptr<EngineHelper> helper,
-               RecvPostCallback recv_post_callback)
+               RecvPostCallback recv_post_callback,
+               ResolveAddrCallback resolve_addr_callback)
     : config_(config),
       self_(self),
       stop_(false),
       next_handle_(util::Random<Handle::ValueType>(bitgen_)),
       next_reqid_(util::Random<ReqId::ValueType>(bitgen_)),
       outgoing_(helper->Metrics(), helper->Tracer()),
       incoming_(helper->Metrics(), helper->Tracer(),
-                std::move(recv_post_callback)),
+                std::move(recv_post_callback),
+                std::move(resolve_addr_callback)),
       metrics_(helper->Metrics()),
       tracer_(helper->Tracer()),
       helper_(std::move(helper)) {
@@ -291,7 +295,9 @@ void Engine::processWrite(Workers& workers, const Handle handle,
     chunk.reqid = req_info.reqid;
     chunk.index = chunk_t(i);
     chunk.byte_offset = offset;
-    chunk.raddr = addr_t(req_info.raddr.value() + offset);
+    chunk.raddr = req_info.raddr.value() == 0
+                      ? addr_t(0)
+                      : addr_t(req_info.raddr.value() + offset);
     if ABSL_PREDICT_FALSE (trace_ctx.is_traced) {
       chunk.trace.SetTraced();
     }

peregrine/src/internal/engine/engine.h

@@ -68,7 +68,8 @@ class Engine final {
   // Creates an engine.
   static std::unique_ptr<Engine> Create(
       const Config& config, HostInfo& self, Control& control,
-      RecvPostCallback recv_post_callback = nullptr);
+      RecvPostCallback recv_post_callback = nullptr,
+      ResolveAddrCallback resolve_addr_callback = nullptr);
 
   // Disallows copy and move.
   DISALLOW_COPY(Engine);
@@ -102,7 +103,8 @@ class Engine final {
   // Constructor.
   Engine(const Config& config, const HostInfo& self,
          std::unique_ptr<EngineHelper> helper,
-         RecvPostCallback recv_post_callback);
+         RecvPostCallback recv_post_callback,
+         ResolveAddrCallback resolve_addr_callback);
 
   // Generates a unique handle.
   Handle genHandle() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {

peregrine/src/internal/request/request_tracker.cc

@@ -58,15 +58,30 @@ bool RequestTracker::Add(const Handle handle,
                          const Metadata& metadata) {
   DCHECK(!reqs.empty());
   absl::MutexLock _(trackers_mu_);
-  PostTracker& pt = trackers_[handle];
+  auto [it, inserted] = trackers_.try_emplace(handle);
+  PostTracker& pt = it->second;
   if (!pt.req_map.empty()) return true;
   pt.metadata = metadata;
   pt.pending_reqs = reqs.size();
   pt.req_map.reserve(reqs.size());
   for (const RequestInfo& req : reqs) {
     DCHECK(req.IsValid());
+    addr_t raddr = req.raddr;
+    if (raddr.value() == 0 && resolve_addr_callback_ != nullptr) {
+      raddr = addr_t(
+          reinterpret_cast<uintptr_t>(resolve_addr_callback_(req.metadata)));
+    }
+    if (pt.start_time == PostTracker::kInvalidTime && raddr.value() == 0) {
+      if (inserted) {
+        trackers_.erase(it);
+      } else {
+        pt.req_map.clear();
+        pt.pending_reqs = 0;
+      }
+      return false;
+    }
     pt.req_map[req.reqid] =
-        std::make_unique<ReqTracker>(req.nchunks, req.metadata);
+        std::make_unique<ReqTracker>(req.nchunks, raddr, req.metadata);
   }
   return true;
 }
@@ -138,6 +153,18 @@ void RequestTracker::Update(const Handle handle, const ReqId reqid,
   }
 }
 
+Byte* RequestTracker::GetChunkAddr(const Handle handle, const ReqId reqid,
+                                   const uint64_t byte_offset) const {
+  absl::MutexLock _(trackers_mu_);
+  const auto it = trackers_.find(handle);
+  if ABSL_PREDICT_FALSE (it == trackers_.end()) return nullptr;
+  const auto req_it = it->second.req_map.find(reqid);
+  if ABSL_PREDICT_FALSE (req_it == it->second.req_map.end()) return nullptr;
+  const uintptr_t base = req_it->second->raddr.value();
+  if ABSL_PREDICT_FALSE (base == 0) return nullptr;
+  return reinterpret_cast<Byte*>(base + byte_offset);
+}
+
 ChunkTracker& RequestTracker::Find(const Handle handle, const ReqId reqid) {
   absl::MutexLock _(trackers_mu_);
   return trackers_.at(handle).req_map.at(reqid)->chunks;