[Store] Keep offload-on-evict under HA and stop charging an upsert for the write it takes over

mooncake-store/include/tenant/quota_policy.h

@@ -8,8 +8,8 @@
 // What an object costs is derived from the object itself, never booked by the
 // code that changes it. The key keeps, in its policy words, the bytes the
 // account currently holds for it and what a replaced object still owes, and
-// every write lock released on the key settles the first towards what the
-// object owes now:
+// every write lock released on the key, and every teardown of its object,
+// settles the first towards what the object owes now:
 //
 //   - nothing, once the key holds no object;
 //   - otherwise its size for every memory replica that holds memory: being
@@ -108,6 +108,8 @@ class TenantQuotaPolicy final : public NamespacePolicy {
     friend class test::MasterServiceTestPeer;
 
     static TenantQuotaAccount& AccountOf(const PolicyContext& ctx);
+    // Brings what the account holds for the key to `owed`.
+    static void Settle(const PolicyContext& ctx, uint64_t owed);
 
     StoreControl& store_;
     const Options options_;

mooncake-store/src/master_service.cpp

@@ -349,16 +349,17 @@ struct MasterService::MemoryEviction {
         : policy(service, settings.allow_evict_soft_pinned),
           immediate(service),
           oplog(service),
+          removal(settings.through_oplog
+                      ? static_cast<eviction::Reclaimer&>(oplog)
+                      : static_cast<eviction::Reclaimer&>(immediate)),
           offload_queue(service),
-          offloading(immediate, offload_queue, settings.offload),
+          offloading(removal, offload_queue, settings.offload),
           kv_events(service),
           log(settings.offload.cap),
           evictor(service.namespaces_, policy,
-                  settings.through_oplog
-                      ? static_cast<eviction::Reclaimer&>(oplog)
-                  : settings.offload_on_evict
+                  settings.offload_on_evict
                       ? static_cast<eviction::Reclaimer&>(offloading)
-                      : static_cast<eviction::Reclaimer&>(immediate),
+                      : removal,
                   Observers(settings)) {}
 
     static bool InPlan(const eviction::Plan& plan, const Replica& replica) {
@@ -378,6 +379,10 @@ struct MasterService::MemoryEviction {
     Policy policy;
     ImmediateReclaimer immediate;
     OpLogReclaimer oplog;
+    // How a replica comes off, whatever decided it goes: through the OpLog
+    // under HA, otherwise on the spot. Offload-on-evict wraps it rather than
+    // standing in for it, so HA keeps the offload.
+    eviction::Reclaimer& removal;
     OffloadQueue offload_queue;
     eviction::OffloadingReclaimer offloading;
     KvEventObserver kv_events;
@@ -1608,12 +1613,9 @@ void MasterService::EvictManagedWeightMembersToCold(
     // Not const: ObjectMetadata::IsLeaseExpired takes it by reference.
     auto now = std::chrono::system_clock::now();
     std::shared_lock<std::shared_mutex> shared_lock(master_state_mutex_);
-    // Memory eviction's own reclaimer, without the offload: a member goes cold
-    // only when it holds a readable cold replica already.
-    eviction::Reclaimer& reclaimer =
-        enable_oplog_
-            ? static_cast<eviction::Reclaimer&>(memory_eviction_->oplog)
-            : memory_eviction_->immediate;
+    // Memory eviction's removal, without the offload: a member goes cold only
+    // when it holds a readable cold replica already.
+    eviction::Reclaimer& reclaimer = memory_eviction_->removal;
     for (const std::string& member_key :
          ns->objects.GroupMembers(revision.manifest.payload_group_id)) {
         if (!selected_members.contains(member_key)) {

mooncake-store/src/tenant/quota_policy.cpp

@@ -216,13 +216,15 @@ void TenantQuotaPolicy::OnReplace(const PolicyContext& ctx) {
 
 void TenantQuotaPolicy::OnTearDown(const PolicyContext& ctx) {
     // The object goes for good, and with it any charge it carried for a
-    // replaced one.
+    // replaced one. The key owes nothing from here, so the account lets go of
+    // it now rather than when the lock is: a new write of the key under this
+    // same lock, such as an upsert taking over a write that completed
+    // nothing, is not charged on top of the object it took over.
     Carried(ctx) = 0;
+    Settle(ctx, 0);
 }
 
 void TenantQuotaPolicy::OnWriteRelease(const PolicyContext& ctx) {
-    TenantQuotaAccount& account = AccountOf(ctx);
-    const uint64_t held = Held(ctx);
     uint64_t& carried = Carried(ctx);
     uint64_t owed = 0;
     if (ctx.guard.has_object()) {
@@ -234,6 +236,18 @@ void TenantQuotaPolicy::OnWriteRelease(const PolicyContext& ctx) {
     } else {
         carried = 0;
     }
+    Settle(ctx, owed);
+    // Ends only now, with what the write leaves settled and its object, if
+    // any, published.
+    if (WriteBegun(ctx)) {
+        SetWriteBegun(ctx, false);
+        AccountOf(ctx).EndWrite();
+    }
+}
+
+void TenantQuotaPolicy::Settle(const PolicyContext& ctx, uint64_t owed) {
+    TenantQuotaAccount& account = AccountOf(ctx);
+    const uint64_t held = Held(ctx);
     if (owed > held) {
         // Growth the core did not ask for up front, such as a write that
         // completed more than it reserved: charged regardless of the quota,
@@ -247,12 +261,6 @@ void TenantQuotaPolicy::OnWriteRelease(const PolicyContext& ctx) {
         }
     }
     SetHeld(ctx, owed);
-    // Ends only now, with what the write leaves settled and its object, if
-    // any, published.
-    if (WriteBegun(ctx)) {
-        SetWriteBegun(ctx, false);
-        account.EndWrite();
-    }
 }
 
 void TenantQuotaPolicy::OnRestored() {

mooncake-store/tests/ha/master_service_ha_test.cpp

@@ -7139,6 +7139,64 @@ TEST_F(MasterServiceHATest, BatchEvictCommitFailureRestoresRemovedReplicas) {
     EXPECT_EQ(ErrorCode::TASK_PENDING_LIMIT_EXCEEDED, full.error());
 }
 
+TEST_F(MasterServiceHATest, OffloadOnEvictOffloadsBeforeEvictingThroughOpLog) {
+    const std::string cluster_id = "test_offload_on_evict_ha_cluster";
+    auto backend = std::make_shared<FakeBatchHaKvBackend>();
+    auto service_config = MasterServiceConfig::builder()
+                              .set_default_kv_lease_ttl(50)
+                              .set_enable_ha(true)
+                              .set_enable_oplog(true)
+                              .set_cluster_id(cluster_id)
+                              .set_oplog_batch_max_entries(1)
+                              .set_enable_offload(true)
+                              .build();
+    service_config.offload_on_evict = true;
+    MasterService service(service_config);
+    ASSERT_EQ(
+        ErrorCode::OK,
+        MasterServiceTestPeer(service).SetBatchOpLogBackendForTesting(backend));
+
+    const std::string segment_name = "offload_on_evict_ha_segment";
+    auto mounted = PrepareSimpleSegment(service, segment_name);
+    ASSERT_TRUE(service
+                    .MountLocalDiskSegment(mounted.client_id,
+                                           /*enable_offloading=*/true)
+                    .has_value());
+    const std::string key = "offload_on_evict_ha_key";
+    PutObjectOnSegment(service, mounted.client_id, key, segment_name);
+
+    // The only memory replica stays until its offload lands.
+    std::this_thread::sleep_for(std::chrono::milliseconds(60));
+    auto evicted = MasterServiceTestPeer(service).EvictNamespaceMemory(
+        kDefaultTenant, 1024);
+    EXPECT_EQ(0u, evicted.freed_bytes);
+    auto queued = service.OffloadObjectHeartbeat(mounted.client_id, true);
+    ASSERT_TRUE(queued.has_value()) << toString(queued.error());
+    ASSERT_EQ(1u, queued->size());
+    EXPECT_EQ(key, queued->front().key);
+    auto before_offload = service.GetReplicaList(key, kDefaultTenant);
+    ASSERT_TRUE(before_offload.has_value());
+    ASSERT_EQ(1u, before_offload->replicas.size());
+    EXPECT_TRUE(before_offload->replicas.front().is_memory_replica());
+
+    // With the offload landed, the next round takes the memory replica off
+    // through the OpLog.
+    StorageObjectMetadata offloaded;
+    offloaded.data_size = 1024;
+    offloaded.transport_endpoint = "offload_on_evict_ha_disk";
+    ASSERT_TRUE(
+        service.NotifyOffloadSuccess(mounted.client_id, *queued, {offloaded})
+            .has_value());
+    std::this_thread::sleep_for(std::chrono::milliseconds(60));
+    evicted = MasterServiceTestPeer(service).EvictNamespaceMemory(
+        kDefaultTenant, 1024);
+    EXPECT_EQ(1024u, evicted.freed_bytes);
+    auto after_evict = service.GetReplicaList(key, kDefaultTenant);
+    ASSERT_TRUE(after_evict.has_value());
+    ASSERT_EQ(1u, after_evict->replicas.size());
+    EXPECT_TRUE(after_evict->replicas.front().is_local_disk_replica());
+}
+
 TEST_F(MasterServiceHATest, PutStartExpiredOverwriteWritesBatchRecordOpLog) {
     const std::string cluster_id =
         "test_batch_record_put_start_cleanup_cluster";

mooncake-store/tests/master_service_tenant_quota_test.cpp

@@ -703,6 +703,30 @@ TEST_F(MasterServiceTenantQuotaTest, PutRevokeRefundsStartCharge) {
     EXPECT_EQ(Snapshot(service, TenantId("tenant-a")).charged_bytes, 0);
 }
 
+TEST_F(MasterServiceTenantQuotaTest,
+       UpsertTakingOverAnUnfinishedWriteIsNotChargedForIt) {
+    const TenantId tenant_id("tenant-a");
+    MasterService service(MakeConfig({{tenant_id, 100}}));
+    UUID client_id = MountSegment(service);
+
+    auto start =
+        service.PutStart(client_id, "key", tenant_id, 100, MemoryConfig());
+    ASSERT_TRUE(start.has_value()) << toString(start.error());
+    ASSERT_EQ(Snapshot(service, tenant_id).charged_bytes, 100);
+
+    // The write it takes over completed nothing, so only the upsert's own
+    // write is charged.
+    auto upsert =
+        service.UpsertStart(client_id, "key", tenant_id, 100, MemoryConfig());
+    ASSERT_TRUE(upsert.has_value()) << toString(upsert.error());
+    EXPECT_EQ(Snapshot(service, tenant_id).charged_bytes, 100);
+
+    auto end =
+        service.UpsertEnd(client_id, "key", tenant_id, ReplicaType::MEMORY);
+    ASSERT_TRUE(end.has_value()) << toString(end.error());
+    EXPECT_EQ(Snapshot(service, tenant_id).charged_bytes, 100);
+}
+
 TEST_F(MasterServiceTenantQuotaTest,
        SizeChangingUpsertTransfersAndReleasesReplacementCharge) {
     const TenantId tenant_id("tenant-a");

mooncake-store/tests/tenant_quota_policy_test.cpp

@@ -371,6 +371,30 @@ TEST_F(TenantQuotaPolicyTest, TearingDownAReplacementDropsWhatItCarried) {
     EXPECT_TRUE(ns.Empty());
 }
 
+TEST_F(TenantQuotaPolicyTest, ATornDownObjectIsNotChargedAgainstItsSuccessor) {
+    Register(kTenantA, 100);
+    // A write that completed nothing.
+    Put(kTenantA, "k", 100, 100, {ReplicaStatus::PROCESSING});
+    ASSERT_EQ(Charged(kTenantA), 100u);
+    auto& ns = NamespaceOf(kTenantA);
+
+    // An upsert takes the write over under one lock: the old object goes
+    // first, so the new one fits in the quota the old one held.
+    {
+        auto guard = ns.objects.Write("k");
+        ASSERT_TRUE(guard.has_value());
+        ASSERT_EQ(guard->metadata().EraseReplicas(
+                      [](const Replica&) { return true; }),
+                  1u);
+        ns.BeforeTearDown(*guard);
+        guard->TearDown();
+        ASSERT_TRUE(ns.Grow(*guard, 100).has_value());
+        PublishUnder(*guard, 100, MemoryReplicas({ReplicaStatus::PROCESSING}));
+    }
+    EXPECT_EQ(Charged(kTenantA), 100u);
+    EXPECT_EQ(Words(kTenantA, "k"), (std::array<uint64_t, 2>{100, 0}));
+}
+
 TEST_F(TenantQuotaPolicyTest, RestoreRebuildsAccountsAndKeyWords) {
     Register(kTenantA, 1000);
     Register(kTenantB, 1000);