perf: reuse dense sorted-set pages during point mutations

src/storage/engine/grouped/zset.cpp

@@ -1028,13 +1028,28 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
   };
   // A serial scan retains one physical read lease and borrowed member views.
   // Keep scratch admitted until both disappear; replies own separate copies.
+  const bool point_write =
+      ordered_probe == nullptr &&
+      ((operation.kind_ == SortedSetOperationKind::kAdd &&
+        operation.entries_.size() == 1) ||
+       (operation.kind_ == SortedSetOperationKind::kRemove &&
+        operation.members_.size() == 1));
   SortedSetOrderedProbe local_probe{.source_ = object, .pages_ = {}};
   if (ordered_probe == nullptr) ordered_probe = &local_probe;
   if (ordered_probe->source_.get() != object.get())
     co_return absl::DataLossError("Sorted Set probe logical view mismatch");
   auto& probed_pages = ordered_probe->pages_;
   using OwnedScanPage = SortedSetOrderedProbe::Page;
   auto retain_probe = [&](std::size_t i, OwnedScanPage page) {
+    if (point_write &&
+        metadata[i].encoded_bytes_ <= kCollectionGroupTargetBytes) {
+      // A point writer needs at most its source and destination boundary.
+      // Dense ordinary pages can carry much more admission than encoded
+      // bytes: decoder headroom is charged per member and twice per read.
+      // Keep that full protection, but do not mistake it for page size.
+      if (probed_pages.size() < 2) probed_pages.emplace(i, std::move(page));
+      return;
+    }
     // Reuse is optional, not a batch-sized payload cache. Bound admitted
     // retention (including decoded entry overhead), so large members and
     // wide pops keep the ordinary one-page scan/fallback under maxmemory.
@@ -1047,6 +1062,37 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
     }
     probed_pages.emplace(i, std::move(page));
   };
+#if LAVIK_FAULTS_ENABLED
+  bool probe_reservation_failure_injected = false;
+#endif
+  auto reserve_work =
+      [&](const GroupedScratchBudget& budget,
+          std::size_t copies) -> absl::StatusOr<MemoryReservation> {
+    if (point_write && !probed_pages.empty()) {
+      if (copies == 0 ||
+          budget.bytes() > std::numeric_limits<std::size_t>::max() / copies)
+        return absl::ResourceExhaustedError(
+            "grouped scratch peak size overflow");
+      bool retain_pages = true;
+      LAVIK_FAULT_INJECT(
+          if (!probe_reservation_failure_injected &&
+              LAVIK_FAULT_MATCHES("LAVIK_FAIL_ZSET_PROBE_RETAIN_RESERVE_KEY",
+                                  key)) {
+            probe_reservation_failure_injected = true;
+            retain_pages = false;
+          });
+      auto reservation = retain_pages
+                             ? TryReserveMemory(budget.bytes() * copies)
+                             : std::nullopt;
+      if (reservation) return std::move(*reservation);
+      // Retention must not turn a formerly admissible point edit into OOM.
+      // Drop the optional owned pages and their credits together, then use
+      // the ordinary read/plan path. Only that final attempt records a
+      // rejection; a failed cache-preserving attempt is not a failed command.
+      probed_pages.clear();
+    }
+    return budget.Reserve(copies);
+  };
   SortedSetPageWorkBudget page_work;
   auto read_page_impl = [&]<typename Page>(std::size_t i,
                                            GroupReadLimiter* limiter = nullptr)
@@ -1092,7 +1138,7 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
     // A parallel wave admits all owned pages before launching children. Its
     // reservation outlives the returned pages; serial probes own theirs here.
     if (limiter == nullptr) {
-      auto reserved = budget.Reserve(2);
+      auto reserved = reserve_work(budget, 2);
       if (!reserved.ok()) co_return reserved.status();
       admission = std::move(*reserved);
     }
@@ -1878,7 +1924,7 @@ StorageEngine::Impl::ExecuteGroupedSortedSetLocked(
     status = budget.AddBytes(member.size() + 256);
     if (!status.ok()) co_return status;
   }
-  auto admission = budget.Reserve(4);
+  auto admission = reserve_work(budget, 4);
   if (!admission.ok()) co_return admission.status();
   working_admission.emplace(std::move(*admission));
   auto check_plan_read = [&]([[maybe_unused]] std::size_t i) -> absl::Status {

tests/grouped/zset_write_e2e_test.cpp

@@ -172,6 +172,253 @@ TEST(GroupedSortedSetWriteE2e, SourceAndPopPagesAvoidSecondPlanRead) {
   EXPECT_EQ(client.Command({"ZRANK", key, member(12)}).text_, "10");
 }
 
+std::string DenseZSetMember(unsigned index) {
+  auto member = std::to_string(index);
+  member.insert(0, 31 - member.size(), '0');
+  // One later-score member sorts before the first page's upper boundary.
+  // Moving it onto that score requires the destination boundary probe.
+  member.insert(member.begin(), index == 500 ? 'a' : 'm');
+  return member;
+}
+
+std::vector<std::string> DenseZSetSeed(
+    std::string_view key, std::map<std::string, double>* expected) {
+  std::vector<std::string> command{"ZADD", std::string(key)};
+  for (unsigned index = 0; index < 1024; ++index) {
+    const auto member = DenseZSetMember(index);
+    command.insert(command.end(), {std::to_string(index), member});
+    expected->emplace(member, index);
+  }
+  return command;
+}
+
+void ExpectDenseZSet(Client& client, std::string_view key,
+                     const std::map<std::string, double>& expected) {
+  EXPECT_EQ(client.Command({"ZCARD", std::string(key)}).text_,
+            std::to_string(expected.size()));
+  std::vector<std::string> lookup{"ZMSCORE", std::string(key)};
+  std::vector<std::pair<double, std::string>> ordered;
+  for (const auto& [member, score] : expected) {
+    lookup.push_back(member);
+    ordered.emplace_back(score, member);
+  }
+  const auto scores = client.Command(lookup);
+  ASSERT_EQ(scores.items_.size(), expected.size());
+  std::size_t index = 0;
+  for (const auto& [member, score] : expected) {
+    ASSERT_EQ(scores.items_[index].kind_, '$') << member;
+    EXPECT_EQ(std::stod(scores.items_[index++].text_), score) << member;
+  }
+  std::sort(ordered.begin(), ordered.end());
+  const auto range =
+      client.Command({"ZRANGE", std::string(key), "0", "-1", "WITHSCORES"});
+  ASSERT_EQ(range.items_.size(), ordered.size() * 2);
+  for (index = 0; index < ordered.size(); ++index) {
+    EXPECT_EQ(range.items_[2 * index].text_, ordered[index].second);
+    EXPECT_EQ(std::stod(range.items_[2 * index + 1].text_),
+              ordered[index].first);
+  }
+}
+
+std::optional<OrderedGroupSnapshot> FirstDenseOrderedPage(
+    const PrivateDisk& disk, std::string_view key) {
+  std::optional<OrderedGroupSnapshot> first;
+  std::uint64_t sequence = 0;
+  disk.VisitRecords([&](const BlockHeader&, const RecordHeader& record,
+                        std::string_view name, std::string_view payload) {
+    if (name != key || !record.auxiliary_group_ || record.external_ ||
+        record.group_prefix_bits_ != 0)
+      return;
+    const auto page = DecodeOrderedGroup(payload);
+    if (!page.ok() || page->previous_ != 0 || page->retired_ ||
+        record.mutation_sequence_ < sequence)
+      return;
+    first = *page;
+    sequence = record.mutation_sequence_;
+  });
+  return first;
+}
+
+TEST(GroupedSortedSetWriteE2e, DensePointSourceAvoidsSecondPlanRead) {
+#if !LAVIK_TEST_FAULTS_AVAILABLE
+  GTEST_SKIP() << "requires ordered plan read failure injection";
+#endif
+  PrivateDisk disk;
+  const std::string key = "dense-source-probe";
+  std::map<std::string, double> expected;
+  {
+    Server server(disk, 1, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_EQ(client.Command(DenseZSetSeed(key, &expected)).text_, "1024");
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  ASSERT_EQ(disk.LatestRootGrouped(key), true);
+  {
+    ScopedEnvironment fault("LAVIK_FAIL_ZSET_PLAN_READ_KEY", key.c_str());
+    ScopedEnvironment page("LAVIK_FAIL_ZSET_PLAN_READ_PAGE", "1");
+    Server server(disk, 3, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_NE(client.Command({"INFO", "MEMORY"})
+                  .text_.find("maxmemory:134217728\r\n"),
+              std::string::npos);
+    // Dense 32-byte members have small encoded pages but substantial decoder
+    // admission. A recovered source page must survive into the mutation plan
+    // with its full charge, both for in-page edits and moves to another page.
+    ASSERT_EQ(
+        client.Command({"ZINCRBY", key, "0.5", DenseZSetMember(10)}).text_,
+        "10.5");
+    expected[DenseZSetMember(10)] = 10.5;
+    ASSERT_EQ(
+        client.Command({"ZINCRBY", key, "2000", DenseZSetMember(11)}).text_,
+        "2011");
+    expected[DenseZSetMember(11)] = 2011;
+    ASSERT_EQ(client.Command({"ZREM", key, DenseZSetMember(12)}).text_, "1");
+    expected.erase(DenseZSetMember(12));
+    // The fault is live: an insertion with no source/boundary probe must
+    // still perform and fail its first planner read, without changing either
+    // graph or hiding the failure behind the cache.
+    const auto failed = client.Command({"ZADD", key, "-1", "unprobed"});
+    EXPECT_EQ(failed.kind_, '-');
+    EXPECT_NE(failed.text_.find("plan page read failure"), std::string::npos);
+    ExpectDenseZSet(client, key, expected);
+    EXPECT_EQ(client.Command({"ZSCORE", key, "unprobed"}).text_, "-1");
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  Server recovered(disk, 2, {}, {}, false, 2, "128M");
+  Client client(recovered.port());
+  ExpectDenseZSet(client, key, expected);
+}
+
+TEST(GroupedSortedSetWriteE2e,
+     DensePointDestinationBoundaryAvoidsSecondPlanRead) {
+#if !LAVIK_TEST_FAULTS_AVAILABLE
+  GTEST_SKIP() << "requires ordered plan read failure injection";
+#endif
+  PrivateDisk disk;
+  const std::string key = "dense-boundary-probe";
+  std::map<std::string, double> expected;
+  {
+    Server server(disk, 1, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_EQ(client.Command(DenseZSetSeed(key, &expected)).text_, "1024");
+    // Make room without changing the first page's upper score. The point
+    // move then tests source+boundary reuse, independently of split policy.
+    ASSERT_EQ(client.Command({"ZREM", key, DenseZSetMember(20)}).text_, "1");
+    expected.erase(DenseZSetMember(20));
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  const auto first = FirstDenseOrderedPage(disk, key);
+  ASSERT_TRUE(first.has_value());
+  ASSERT_GT(first->entries_.size(), 128U);
+  const auto boundary_score = first->entries_.back().score_;
+  ASSERT_GT(500, boundary_score);
+  ASSERT_LT(DenseZSetMember(500), first->entries_.back().value_);
+  {
+    ScopedEnvironment fault("LAVIK_FAIL_ZSET_PLAN_READ_KEY", key.c_str());
+    ScopedEnvironment page("LAVIK_FAIL_ZSET_PLAN_READ_PAGE", "1");
+    Server server(disk, 3, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_NE(client.Command({"INFO", "MEMORY"})
+                  .text_.find("maxmemory:134217728\r\n"),
+              std::string::npos);
+    const auto changed =
+        client.Command({"ZINCRBY", key, std::to_string(boundary_score - 500),
+                        DenseZSetMember(500)});
+    ASSERT_EQ(changed.kind_, '$') << changed.text_;
+    EXPECT_EQ(std::stod(changed.text_), boundary_score);
+    expected[DenseZSetMember(500)] = boundary_score;
+    ExpectDenseZSet(client, key, expected);
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  Server recovered(disk, 2, {}, {}, false, 2, "128M");
+  Client client(recovered.port());
+  ExpectDenseZSet(client, key, expected);
+}
+
+TEST(GroupedSortedSetWriteE2e, DensePointProbeRetentionFallsBackWithoutOOM) {
+#if !LAVIK_TEST_FAULTS_AVAILABLE
+  GTEST_SKIP() << "requires probe-retention reservation failure injection";
+#endif
+  PrivateDisk disk;
+  const std::string key = "dense-probe-fallback";
+  std::map<std::string, double> expected;
+  auto rejections = [](Client& client) {
+    const auto info = client.Command({"INFO", "MEMORY"}).text_;
+    const std::string_view field = "oom_rejected_commands:";
+    const auto offset = info.find(field);
+    Check(offset != std::string::npos, "missing OOM rejection counter");
+    return std::stoull(info.substr(offset + field.size()));
+  };
+  {
+    Server server(disk, 1, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_EQ(client.Command(DenseZSetSeed(key, &expected)).text_, "1024");
+    ASSERT_EQ(client.Command({"ZREM", key, DenseZSetMember(20)}).text_, "1");
+    expected.erase(DenseZSetMember(20));
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  const auto first = FirstDenseOrderedPage(disk, key);
+  ASSERT_TRUE(first.has_value());
+  ASSERT_GT(first->entries_.size(), 128U);
+  const auto boundary_score = first->entries_.back().score_;
+  ASSERT_GT(500, boundary_score);
+  ASSERT_LT(DenseZSetMember(500), first->entries_.back().value_);
+  ScopedEnvironment reservation_failure(
+      "LAVIK_FAIL_ZSET_PROBE_RETAIN_RESERVE_KEY", key.c_str());
+  {
+    Server server(disk, 3, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    ASSERT_NE(client.Command({"INFO", "MEMORY"})
+                  .text_.find("maxmemory:134217728\r\n"),
+              std::string::npos);
+    const auto rejected_before = rejections(client);
+    // Only the optional cache-preserving reservation fails. The ordinary
+    // retry still uses the real nonzero memory limit. Same-page edits reach
+    // final-plan admission; a tied destination first needs another probe.
+    ASSERT_EQ(
+        client.Command({"ZINCRBY", key, "0.5", DenseZSetMember(10)}).text_,
+        "10.5");
+    expected[DenseZSetMember(10)] = 10.5;
+    const auto moved =
+        client.Command({"ZINCRBY", key, std::to_string(boundary_score - 500),
+                        DenseZSetMember(500)});
+    ASSERT_EQ(moved.kind_, '$') << moved.text_;
+    EXPECT_EQ(std::stod(moved.text_), boundary_score);
+    expected[DenseZSetMember(500)] = boundary_score;
+    ASSERT_EQ(client.Command({"ZREM", key, DenseZSetMember(12)}).text_, "1");
+    expected.erase(DenseZSetMember(12));
+    ExpectDenseZSet(client, key, expected);
+    EXPECT_EQ(rejections(client), rejected_before);
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  {
+    // Prove the rejection really discards the optional source: its fallback
+    // planner read now fails. Both indexes must retain the previous value.
+    ScopedEnvironment fault("LAVIK_FAIL_ZSET_PLAN_READ_KEY", key.c_str());
+    ScopedEnvironment page("LAVIK_FAIL_ZSET_PLAN_READ_PAGE", "1");
+    Server server(disk, 3, {}, {}, false, 2, "128M");
+    Client client(server.port());
+    const auto rejected_before = rejections(client);
+    const auto failed =
+        client.Command({"ZINCRBY", key, "0.5", DenseZSetMember(10)});
+    EXPECT_EQ(failed.kind_, '-');
+    EXPECT_NE(failed.text_.find("plan page read failure"), std::string::npos);
+    EXPECT_EQ(rejections(client), rejected_before);
+    ExpectDenseZSet(client, key, expected);
+    client.Durable();
+    ASSERT_EQ(server.Wait(true), 0) << server.Log();
+  }
+  Server recovered(disk, 2, {}, {}, false, 2, "128M");
+  Client client(recovered.port());
+  ExpectDenseZSet(client, key, expected);
+}
+
 TEST(GroupedSortedSetWriteE2e,
      PointWritesReuseMemberLeafAndKeepBatchFailureAtomic) {
 #if !LAVIK_TEST_FAULTS_AVAILABLE