[Bugfix][Store] Make local MEMORY preference independent of replica order

mooncake-store/include/replica_selection.h

@@ -117,12 +117,16 @@ inline const Replica::Descriptor *PickBestRemoteMemory(
 
 // Select the best replica from a list: prefer local MEMORY, local NOF_SSD,
 // remote MEMORY, remote NOF_SSD, LOCAL_DISK, DFS, then DISK. Master may return
-// replicas in any order, so we always scan. When scoring is enabled and there
-// are multiple remote MEMORY replicas, the best-scoring one is chosen instead
-// of the first encountered.
+// replicas in any order, so we always scan all replicas before deciding — this
+// ensures local MEMORY is always preferred over local NOF regardless of the
+// order replicas appear in the list. When scoring is enabled and there are
+// multiple remote MEMORY replicas, the best-scoring one is chosen instead of
+// the first encountered.
 inline const Replica::Descriptor *SelectBestReplica(
     const std::vector<Replica::Descriptor> &replicas,
     const std::unordered_set<std::string> &local_endpoints) {
+    const Replica::Descriptor *local_memory = nullptr;
+    const Replica::Descriptor *local_nof = nullptr;
     const Replica::Descriptor *first_memory = nullptr;
     const Replica::Descriptor *first_nof = nullptr;
     for (const auto &r : replicas) {
@@ -131,18 +135,23 @@ inline const Replica::Descriptor *SelectBestReplica(
             if (local_endpoints.count(
                     r.get_memory_descriptor()
                         .buffer_descriptor.transport_endpoint_)) {
-                return &r;  // local MEMORY — best case
+                if (!local_memory) local_memory = &r;
+            } else {
+                if (!first_memory) first_memory = &r;
             }
-            if (!first_memory) first_memory = &r;
         } else if (r.is_nof_replica()) {
             if (local_endpoints.count(
                     r.get_nof_descriptor()
                         .buffer_descriptor.transport_endpoint_)) {
-                return &r;  // local NOF_SSD — also good
+                if (!local_nof) local_nof = &r;
+            } else {
+                if (!first_nof) first_nof = &r;
             }
-            if (!first_nof) first_nof = &r;
         }
     }
+    // Local MEMORY always beats local NOF regardless of list order.
+    if (local_memory) return local_memory;
+    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()) {
@@ -168,4 +177,33 @@ inline const Replica::Descriptor *SelectBestReplica(
     return best;
 }
 
+// Select a complete MEMORY replica for the session range-read path. The
+// session ranged-get path (Client::BatchTransferReadRanges) only ever reads
+// MEMORY replicas: it rejects every non-MEMORY entry, so selecting a NOF_SSD
+// replica here would let a session start and then fail the actual transfer.
+// Among complete MEMORY replicas, a local one is preferred; otherwise the
+// first complete MEMORY replica is returned (nullptr if there is none).
+//
+// Unlike SelectBestReplica this deliberately does NOT consider NOF_SSD: the
+// session path cannot read it. Local MEMORY is still preferred over a remote
+// MEMORY replica regardless of the order the master returned them in.
+inline const Replica::Descriptor *SelectCompleteMemoryReplica(
+    const std::vector<Replica::Descriptor> &replicas,
+    const std::unordered_set<std::string> &local_endpoints) {
+    const Replica::Descriptor *first_memory = nullptr;
+    for (const auto &r : replicas) {
+        if (r.status != ReplicaStatus::COMPLETE || !r.is_memory_replica()) {
+            continue;
+        }
+        if (local_endpoints.count(r.get_memory_descriptor()
+                                      .buffer_descriptor.transport_endpoint_)) {
+            return &r;
+        }
+        if (!first_memory) {
+            first_memory = &r;
+        }
+    }
+    return first_memory;
+}
+
 }  // namespace mooncake

mooncake-store/src/real_client.cpp

@@ -395,6 +395,7 @@ inline tl::expected<void, ErrorCode> gather_maybe_device_to_host(
 // SelectBestReplica and the replica-scoring helpers live in
 // replica_selection.h (included above) so they can be unit-tested directly.
 using mooncake::SelectBestReplica;
+using mooncake::SelectCompleteMemoryReplica;
 
 // Build a QueryResult containing only the chosen replica so that
 // Client::Get / Client::BatchGet (which internally call
@@ -413,25 +414,6 @@ inline bool is_object_range_overflow(size_t offset, size_t size, size_t limit) {
     return size > limit || offset > limit - size;
 }
 
-inline const Replica::Descriptor *SelectCompleteMemoryReplica(
-    const std::vector<Replica::Descriptor> &replicas,
-    const std::unordered_set<std::string> &local_endpoints) {
-    const Replica::Descriptor *first_memory = nullptr;
-    for (const auto &r : replicas) {
-        if (r.status != ReplicaStatus::COMPLETE || !r.is_memory_replica()) {
-            continue;
-        }
-        if (local_endpoints.count(r.get_memory_descriptor()
-                                      .buffer_descriptor.transport_endpoint_)) {
-            return &r;
-        }
-        if (!first_memory) {
-            first_memory = &r;
-        }
-    }
-    return first_memory;
-}
-
 inline bool HasMemoryReplica(const std::vector<Replica::Descriptor> &replicas) {
     for (const auto &r : replicas) {
         if (r.is_memory_replica()) {
@@ -5478,6 +5460,11 @@ std::vector<int> RealClient::batch_get_session_start(
             continue;
         }
 
+        // The session range-read path (batch_get_into_multi_buffer_ranges ->
+        // Client::BatchTransferReadRanges) only reads MEMORY replicas, so the
+        // session must be started against a complete MEMORY replica. Selecting
+        // a NOF_SSD replica here would let the session start and then fail the
+        // actual transfer (BatchTransferReadRanges rejects non-MEMORY entries).
         const auto *replica =
             SelectCompleteMemoryReplica(query_result.replicas, local_endpoints);
         if (!replica) {
@@ -5539,8 +5526,10 @@ std::vector<int> RealClient::batch_get_into_multi_buffer_ranges(
                 results[i] = static_cast<int>(toInt(ErrorCode::LEASE_EXPIRED));
                 continue;
             }
-            // start cached a single complete memory replica via
-            // FilterQueryResult.
+            // start cached a single complete MEMORY replica via
+            // FilterQueryResult. The session range-read path only reads
+            // MEMORY replicas, so the cached replica is always MEMORY here;
+            // extract its object size for overflow detection.
             const auto &replica = it->second.replicas.front();
             const size_t replica_limit =
                 replica.is_memory_replica()

mooncake-store/tests/replica_selection_test.cpp

@@ -97,6 +97,74 @@ TEST_F(ReplicaSelectionTest, LocalMemoryAlwaysWins) {
         "nodeB");  // locality beats protocol
 }
 
+// Regression for issue #3658: a complete local MEMORY replica must win over a
+// complete local NOF replica regardless of the order they appear in the list.
+// ObjectMetadata preserves insertion order, and promotion appends a new MEMORY
+// replica to existing metadata, so a MEMORY replica may legally appear after a
+// NOF replica. The session range-read path is memory-only, so selecting NOF
+// here would regress a session that still has a usable MEMORY replica.
+TEST_F(ReplicaSelectionTest, LocalMemoryBeatsLocalNoFInBothOrders) {
+    std::unordered_set<std::string> local = {"memB", "nofB"};
+
+    // local NOF listed before local MEMORY
+    std::vector<Replica::Descriptor> nof_first = {
+        MakeNoF("nofB"),
+        MakeMemory("memB", "tcp"),
+    };
+    const auto* sel = SelectBestReplica(nof_first, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_memory_replica());
+    EXPECT_EQ(
+        sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_,
+        "memB");
+
+    // local MEMORY listed before local NOF (the previously-correct order)
+    std::vector<Replica::Descriptor> mem_first = {
+        MakeMemory("memB", "tcp"),
+        MakeNoF("nofB"),
+    };
+    sel = SelectBestReplica(mem_first, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_memory_replica());
+    EXPECT_EQ(
+        sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_,
+        "memB");
+}
+
+// Regression for issue #3658 (session range-read path): the session ranged-get
+// path only reads MEMORY replicas, so a session must never be started against a
+// NOF_SSD replica — it would pass the size check and then fail the actual
+// transfer. SelectCompleteMemoryReplica must therefore skip a NOF-only object
+// and return nullptr (the caller then rejects the session), even though
+// SelectBestReplica would have selected the NOF replica for the single-object
+// get path.
+TEST_F(ReplicaSelectionTest, SessionPathRejectsNoFOnlyObject) {
+    std::unordered_set<std::string> local = {"nofB"};
+
+    // A NOF-only object: SelectBestReplica returns the NOF replica (it is the
+    // last-resort tier for the single-object get path), but the session path
+    // cannot read NOF, so the session selector must decline it.
+    std::vector<Replica::Descriptor> nof_only = {MakeNoF("nofB")};
+    const auto* best = SelectBestReplica(nof_only, local);
+    ASSERT_NE(best, nullptr);
+    EXPECT_TRUE(best->is_nof_replica());  // single-object path would use it
+    EXPECT_EQ(SelectCompleteMemoryReplica(nof_only, local),
+              nullptr);  // session path: nothing to read
+
+    // A local MEMORY replica must still be selected for the session even when a
+    // local NOF replica is also present (local MEMORY outranks local NOF).
+    std::vector<Replica::Descriptor> mem_and_nof = {
+        MakeNoF("nofB"),
+        MakeMemory("memB", "tcp"),
+    };
+    const auto* sel = SelectCompleteMemoryReplica(mem_and_nof, local);
+    ASSERT_NE(sel, nullptr);
+    EXPECT_TRUE(sel->is_memory_replica());
+    EXPECT_EQ(
+        sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_,
+        "memB");
+}
+
 TEST_F(ReplicaSelectionTest, ScoringOffKeepsFirstRemoteMemory) {
     // No scorer injected and env not set -> must return the FIRST remote
     // MEMORY.