MB-69674: Yield/Cancel when over the backfill threshold

As the CacheTransfer task can raise memory usage.

all_keys:true yield (don't even visit) if we're over the backfill
threshold. This would mean a value-eviction rebalance is delayed,
but similarly any DCP rebalance which is backfilling under high
memory would be delayed until high memory resolves.

all_keys:false cancel the transfer, allowing a rebalance to complete
rather than waiting for high memory situation to resolve. This is
in my view preferable, a rebalance maybe occuring to resolve some
problem, and waiting for memory to reduce just so we can populate
the remote cache to the maximum is of a less priority.

Change-Id: I838d93bd53e74a2f8d832f925c0900b5138214d5
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/238356
Tested-by: Build Bot <build@couchbase.com>
Reviewed-by: Paolo Cocchi <paolo.cocchi@couchbase.com>

Author: jimwwalker

engines/ep/src/dcp/cache_transfer_stream.cc

@@ -252,6 +252,27 @@ bool CacheTransferTask::run() {
         return false;
     }
 
+    // Yield an all_keys transfer when memory usage is over the backfill
+    // threshold, or cancel when this is not an all_keys transfer.
+    if (engine->getKVBucket()->isMemUsageAboveBackfillThreshold()) {
+        if (isAllKeys) {
+            snooze(maybeLogHighMemoryPressure(
+                    visitor.getStream(),
+                    "task yielded due to backfill threshold"));
+            // run again
+            return true;
+        }
+        auto& stream = visitor.getStream();
+        OBJ_LOG_WARN_CTX(stream,
+                         "CacheTransferTask::run: cancelling transfer due "
+                         "to backfill "
+                         "threshold",
+                         {"vb", vbid});
+        stream.cancelTransfer();
+        // stop running the task
+        return false;
+    }
+
     // Ensure we always teardown the hash table visit (drops stream reference)
     auto guard = folly::makeGuard([this] { visitor.tearDownHashTableVisit(); });
 

engines/ep/tests/module_tests/dcp_cache_transfer_test.cc

@@ -88,6 +88,12 @@ class DcpCacheTransferTest : public STParameterizedBucketTest {
                     "for vb:0");
     }
 
+    /**
+     * Helper method to test backfill threshold behavior.
+     * @param allKeys Whether to use all_keys mode
+     */
+    void testBackfillThresholdBehavior(bool allKeys);
+
     std::unordered_set<Item> expectedItems;
 };
 
@@ -656,6 +662,89 @@ TEST_P(DcpCacheTransferTest, key_only_transfer) {
     EXPECT_EQ(cb::mcbp::DcpStreamEndStatus::Ok, producers.last_end_status);
 }
 
+void DcpCacheTransferTest::testBackfillThresholdBehavior(bool allKeys) {
+    // This test relies on precise memory tracking
+    const auto& stats = engine->getEpStats();
+    if (!stats.isMemoryTrackingEnabled()) {
+        GTEST_SKIP();
+    }
+
+    // Need static clock for this test so we can run the yielded task
+    cb::time::steady_clock::use_chrono = false;
+    auto scopeGuard = folly::makeGuard(
+            []() { cb::time::steady_clock::use_chrono = true; });
+
+    // Store items for testing
+    expectedItems.insert(store_item(Vbid(0), makeStoredDocKey("k2"), "2"));
+
+    // Remember original max data size
+    const auto originalMaxDataSize = stats.getMaxDataSize();
+
+    // Set bucket quota using current memory usage such that we exceed the
+    // backfill threshold
+    double backfillThreshold =
+            engine->getConfiguration().getBackfillMemThreshold() / 100.0;
+    engine->setMaxDataSize(stats.getPreciseTotalMemoryUsed() /
+                           backfillThreshold);
+
+    // Verify we're actually above the threshold
+    ASSERT_TRUE(store->isMemUsageAboveBackfillThreshold());
+
+    // Create stream with or without all_keys based on parameter
+    auto stream = createStream(
+            *producer,
+            1,
+            Vbid(0),
+            store->getVBucket(vbid)->getHighSeqno(),
+            store->getVBucket(vbid)->getHighSeqno(),
+            nlohmann::json{{"cts", {{"all_keys", allKeys}}}}.dump());
+
+    // Run the cache transfer task
+    runCacheTransferTask();
+
+    if (allKeys) {
+        // Yield case: stream should remain active, no items queued yet
+        EXPECT_TRUE(stream->isActive());
+        EXPECT_EQ(0, stream->getItemsRemaining())
+                << "Task should yield without queueing items or stream-end";
+
+        // Restore memory to normal
+        engine->setMaxDataSize(originalMaxDataSize);
+        ASSERT_FALSE(store->isMemUsageAboveBackfillThreshold());
+
+        // Advance time so the yielded task can run again
+        cb::time::steady_clock::advance(std::chrono::seconds(1));
+
+        // Now the transfer should complete successfully
+        runCacheTransferTask();
+        ASSERT_EQ(3, stream->getItemsRemaining());
+        EXPECT_TRUE(stream->validateNextResponse(expectedItems));
+        EXPECT_TRUE(stream->validateNextResponse(expectedItems));
+        EXPECT_TRUE(stream->validateNextResponseIsEnd());
+    } else {
+        // Cancel case: stream should be cancelled with stream-end queued
+        EXPECT_FALSE(stream->isActive());
+        ASSERT_EQ(1, stream->getItemsRemaining())
+                << "Only stream-end should be queued, no items transferred";
+        EXPECT_TRUE(stream->validateNextResponseIsEnd());
+
+        // Restore original memory limit
+        engine->setMaxDataSize(originalMaxDataSize);
+    }
+}
+
+// Test that CacheTransferTask yields when all_keys=true and memory usage is
+// above backfill threshold
+TEST_P(DcpCacheTransferTest, backfill_threshold_all_keys_yields) {
+    testBackfillThresholdBehavior(true);
+}
+
+// Test that CacheTransferTask cancels when all_keys=false and memory usage is
+// above backfill threshold
+TEST_P(DcpCacheTransferTest, backfill_threshold_cancels) {
+    testBackfillThresholdBehavior(false);
+}
+
 // We only test persistent buckets. Ephemeral doesn't apply nor will it work as
 // we cannot transfer the linked list or system events...
 INSTANTIATE_TEST_SUITE_P(