From ce7e97c23a209879ab560bc90103b73a92de47d7 Mon Sep 17 00:00:00 2001
From: 2sumtech <2sumtech@gmail.com>
Date: Thu, 27 Aug 2026 10:00:09 -0700
Subject: [PATCH] [Bugfix][Store] Keep scanning for a local MEMORY replica past
 a local NoF one

SelectBestReplica documents the priority local MEMORY > local NOF_SSD >
remote MEMORY > remote NOF_SSD and notes that the master may return
replicas in any order, so the scan must not stop early. It did stop
early: the first COMPLETE local NOF_SSD replica was returned immediately,
so a local MEMORY replica listed after it was never seen and the read was
served from NVMe-over-Fabrics instead of local host memory.

Remember the first local NOF_SSD replica and keep scanning instead. It is
returned after the loop, so it still outranks every remote replica.
---
 mooncake-store/include/replica_selection.h    |  8 +++-
 .../tests/replica_selection_test.cpp          | 38 +++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/mooncake-store/include/replica_selection.h b/mooncake-store/include/replica_selection.h
index eb7fef190b..125f9cad7b 100644
--- a/mooncake-store/include/replica_selection.h
+++ b/mooncake-store/include/replica_selection.h
@@ -125,6 +125,7 @@ inline const Replica::Descriptor *SelectBestReplica(
     const std::unordered_set<std::string> &local_endpoints) {
     const Replica::Descriptor *first_memory = nullptr;
     const Replica::Descriptor *first_nof = nullptr;
+    const Replica::Descriptor *local_nof = nullptr;
     for (const auto &r : replicas) {
         if (r.status != ReplicaStatus::COMPLETE) continue;
         if (r.is_memory_replica()) {
@@ -138,11 +139,16 @@ inline const Replica::Descriptor *SelectBestReplica(
             if (local_endpoints.count(
                     r.get_nof_descriptor()
                         .buffer_descriptor.transport_endpoint_)) {
-                return &r;  // local NOF_SSD — also good
+                // local NOF_SSD — good, but local MEMORY outranks it and may
+                // still appear later, so remember it and keep scanning.
+                if (!local_nof) local_nof = &r;
+                continue;
             }
             if (!first_nof) first_nof = &r;
         }
     }
+    // No local MEMORY replica: a local NOF_SSD still beats every remote one.
+    if (local_nof) return local_nof;
     // No local replica. Among remote MEMORY replicas, optionally pick the
     // best-scoring one instead of the first encountered (issue #2516).
     if (first_memory && RemoteReplicaScoringEnabled()) {
diff --git a/mooncake-store/tests/replica_selection_test.cpp b/mooncake-store/tests/replica_selection_test.cpp
index af5600bad7..c571da1197 100644
--- a/mooncake-store/tests/replica_selection_test.cpp
+++ b/mooncake-store/tests/replica_selection_test.cpp
@@ -97,6 +97,44 @@ TEST_F(ReplicaSelectionTest, LocalMemoryAlwaysWins) {
         "nodeB");  // locality beats protocol
 }
 
+TEST_F(ReplicaSelectionTest, LocalMemoryWinsOverEarlierLocalNoF) {
+    // Master may return replicas in any order, so a local NOF_SSD listed
+    // before a local MEMORY replica must not win: local MEMORY outranks it.
+    std::unordered_set<std::string> local = {"nodeA"};
+    std::vector<Replica::Descriptor> reps = {
+        MakeNoF("nodeA"),             // local, but slower tier
+        MakeMemory("nodeA", "rdma"),  // local MEMORY -> must win
+    };
+    const auto* sel = SelectBestReplica(reps, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_memory_replica());
+}
+
+TEST_F(ReplicaSelectionTest, IncompleteLocalMemoryFallsBackToLocalNoF) {
+    // The local MEMORY replica is not readable yet, so the local NOF_SSD one
+    // remains the best choice.
+    std::unordered_set<std::string> local = {"nodeA"};
+    std::vector<Replica::Descriptor> reps = {
+        MakeNoF("nodeA"),
+        MakeMemory("nodeA", "rdma", ReplicaStatus::PROCESSING),
+    };
+    const auto* sel = SelectBestReplica(reps, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_nof_replica());
+}
+
+TEST_F(ReplicaSelectionTest, LocalNoFPrecedesRemoteMemory) {
+    // Locality still outranks tier for remote MEMORY replicas.
+    std::unordered_set<std::string> local = {"nodeA"};
+    std::vector<Replica::Descriptor> reps = {
+        MakeNoF("nodeA"),
+        MakeMemory("nodeB", "rdma"),
+    };
+    const auto* sel = SelectBestReplica(reps, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_nof_replica());
+}
+
 TEST_F(ReplicaSelectionTest, ScoringOffKeepsFirstRemoteMemory) {
     // No scorer injected and env not set -> must return the FIRST remote
     // MEMORY.