perf(storage): avoid duplicate member reads in ZRANK

src/storage/engine/grouped/zset.cpp

@@ -1470,9 +1470,11 @@ StorageEngine::Impl::TryConcurrentSortedSetPageRead(
       return ReadGroupedSortedSetRandom(store, partition, db_id, key, digest,
                                         operation, current, random.draws_,
                                         std::move(random_result), true);
-    return ExecuteGroupedSortedSetLocked(store, partition, db_id, key, digest,
-                                         operation, current, nullptr, nullptr,
-                                         nullptr, nullptr, nullptr, true);
+    return ExecuteGroupedSortedSetLocked(
+        store, partition, db_id, key, digest, operation, current, nullptr,
+        nullptr, nullptr, nullptr, nullptr, true,
+        operation.kind_ == SortedSetOperationKind::kRank ? &member_score
+                                                         : nullptr);
   }();
   if (!read.ok()) {
     if (EffectiveRecordDbEpoch(partition, db_id) !=
@@ -2712,7 +2714,9 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
     TxShardWrites* tx, ReplicationCommandAppend* replication,
     const MutationPrecondition* mutation_precondition,
     PreparedOrderedMutation* prepared, SortedSetOrderedProbe* ordered_probe,
-    bool group_locked) {
+    bool group_locked, const std::optional<double>* locked_rank_score) {
+  assert(locked_rank_score == nullptr ||
+         (group_locked && operation.kind_ == SortedSetOperationKind::kRank));
   if (!object || !object->is_ordered() ||
       object->ordered_directory().root().kind_ !=
           OrderedCollectionKind::kSortedSet)
@@ -2726,11 +2730,16 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
     co_return co_await ScanGroupedSortedSetMembers(
         store, partition, db_id, key, digest, operation, std::move(object),
         group_locked);
-  auto input_admission = ReserveInputs(operation);
-  if (!input_admission.ok()) co_return input_admission.status();
+  MemoryReservation input_admission;
   Members members;
-  auto status = PrepareMembers(operation, &members);
-  if (!status.ok()) co_return status;
+  absl::Status status;
+  if (locked_rank_score == nullptr) {
+    auto reserved = ReserveInputs(operation);
+    if (!reserved.ok()) co_return reserved.status();
+    input_admission = std::move(*reserved);
+    status = PrepareMembers(operation, &members);
+    if (!status.ok()) co_return status;
+  }
   const auto& directory = object->ordered_directory();
   const auto& metadata = directory.groups();
   SortedSetResult result;
@@ -3518,7 +3527,7 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
   }
   std::size_t remaining_sources = 0;
   std::optional<SortedSetMemberProbe> member_probe;
-  {
+  if (locked_rank_score == nullptr) {
     // Prefix routing retains only per-group metadata. Exact members and
     // scores are decoded from the selected Hash leaves, never trusted from
     // a digest alone. Batch requests read each selected leaf just once.
@@ -3689,7 +3698,11 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
   }
   if (operation.kind_ == SortedSetOperationKind::kRank) {
     const auto member = operation.members_.front();
-    const auto score = members.at(member).before_;
+    // The member S hold spans routing, ordered-lock waits and projection, so
+    // its score (or absence) cannot change. Rank still uses the final common
+    // directory below because unrelated pages may have changed their counts.
+    const auto score =
+        locked_rank_score ? *locked_rank_score : members.at(member).before_;
     if (!score) co_return result;
     // Locate one page by its maximum (score, member), then use Fenwick for
     // the skipped prefix without loading the intervening equal-score pages.

src/storage/engine/impl.h

@@ -2725,15 +2725,17 @@ class StorageEngine::Impl {
   // group_locked requires IS plus S on every accessed member/ordered group;
   // previous is the common logical cut captured after acquiring those holds.
   // Its aggregate counts may outlive unrelated root publications.
+  // locked_rank_score reuses a completed member lookup, including a miss,
+  // while its member S hold remains owned. It must outlive the awaited task.
   Task<absl::StatusOr<SortedSetResult>> ExecuteGroupedSortedSetLocked(
       WorkerStore& store, WorkerStore::PartitionStore& partition,
       std::uint8_t db_id, std::string_view key, const Digest& digest,
       const SortedSetOperation& operation, GroupedObject::Handle previous,
       TxShardWrites* tx, ReplicationCommandAppend* replication,
       const MutationPrecondition* mutation_precondition = nullptr,
       PreparedOrderedMutation* prepared = nullptr,
-      SortedSetOrderedProbe* ordered_probe = nullptr,
-      bool group_locked = false);
+      SortedSetOrderedProbe* ordered_probe = nullptr, bool group_locked = false,
+      const std::optional<double>* locked_rank_score = nullptr);
 
   // Caller holds key IX and database admission. A disengaged result requests
   // one X retry after releasing all groups and IX; errors are never retried.

tests/storage_write_buffer_test.cpp

@@ -31,6 +31,7 @@ class WriteBufferTestPeer {
     kGroupedReadUnmetered,
     kGroupedReadAdmissionFailure,
     kGroupedReadCommonView,
+    kGroupedRankSingleRead,
     kProgress,
     kEpochChanged,
     kSourceChanged,
@@ -898,6 +899,48 @@ class WriteBufferTestPeer {
     if (!object || !object->has_member_index())
       co_return absl::InternalError("fixture did not create member leaves");
 
+    if (scenario == Scenario::kGroupedRankSingleRead) {
+      auto paused = co_await storage.ConfigureDefrag(
+          {.action_ = DefragConfigAction::kPause});
+      if (!paused.ok()) co_return paused;
+      // Drain staging so each member/ordered payload requires physical IO.
+      // Background cleaners are disabled by the fixture; the short parent
+      // key and resident directory require no separate key/root reads.
+      impl.YieldActiveWriteBuffers(store);
+      while (store.flush_running_ || !store.flush_queue_.empty())
+        co_await bycorf::Yield(*store.worker_);
+      for (unsigned index : {0u, 900u, 1024u}) {
+        const bool present = index < names.size();
+        const std::array<std::string_view, 1> requested{
+            present ? std::string_view(names[index]) : "missing"};
+        for (bool reverse : {false, true}) {
+          const auto before =
+              co_await CollectWorkerMetrics(CommandMetricsDetail::kNone);
+          auto read = co_await storage.ExecuteSortedSet(
+              0, key,
+              {.kind_ = SortedSetOperationKind::kRank,
+               .members_ = requested,
+               .reverse_ = reverse});
+          if (!read.ok()) co_return read.status();
+          const auto after =
+              co_await CollectWorkerMetrics(CommandMetricsDetail::kNone);
+          EXPECT_EQ(after.storage_reads_.operations_ -
+                        before.storage_reads_.operations_,
+                    present ? 2 : 1)
+              << "index=" << index << " reverse=" << reverse;
+          EXPECT_EQ(read->length_, 1024);
+          if (present) {
+            EXPECT_EQ(read->rank_, reverse ? 1023 - index : index);
+            EXPECT_EQ(read->rank_score_, index);
+          } else {
+            EXPECT_FALSE(read->rank_);
+            EXPECT_FALSE(read->rank_score_);
+          }
+        }
+      }
+      co_return absl::OkStatus();
+    }
+
     if (scenario == Scenario::kGroupedReadCommonView) {
       unsigned other = 900;
       auto first = object->directory().Find(names[0])->id_;
@@ -1041,7 +1084,8 @@ class WriteBufferTestPeer {
       co_return co_await RunGroupedPublication(storage);
     if (scenario == Scenario::kGroupedReadUnmetered ||
         scenario == Scenario::kGroupedReadAdmissionFailure ||
-        scenario == Scenario::kGroupedReadCommonView)
+        scenario == Scenario::kGroupedReadCommonView ||
+        scenario == Scenario::kGroupedRankSingleRead)
       co_return co_await RunGroupedRead(storage, scenario);
     if (scenario == Scenario::kResumeActiveStreams)
       co_return co_await RunResumeActiveStreams(storage);
@@ -1405,6 +1449,10 @@ TEST(StorageGroupedReadTest, RetainsEarlierLeafWhileWaitingForLaterWriter) {
   RunWriteBufferScenario(WriteBufferTestPeer::Scenario::kGroupedReadCommonView);
 }
 
+TEST(StorageGroupedReadTest, RanksReadEachRequiredPayloadOnce) {
+  RunWriteBufferScenario(WriteBufferTestPeer::Scenario::kGroupedRankSingleRead);
+}
+
 TEST(StorageWriteBufferResumeTest, ThreeStreamsReuseTwoBuffersWithoutSealing) {
   RunWriteBufferScenario(WriteBufferTestPeer::Scenario::kResumeActiveStreams);
 }