[BP] MB-59817 [1/2]: Warmup vbuckets in the order specified

KVBucket::pauseResumeVisit previously only allowed visiting vbuckets
in increasing vbid order, which could cause warmup to load few active
vbuckets.

Change-Id: I6b6a6c09b0f2ef15b60bd5da5c7281140623fb0a
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/201849
Well-Formed: Restriction Checker
Tested-by: Pavlos Georgiou <[email protected]>
Reviewed-by: Dave Rigby <[email protected]>

Author: pavlosg

engines/ep/src/kv_bucket.cc

@@ -2448,6 +2448,21 @@ KVBucket::Position KVBucket::pauseResumeVisit(PauseResumeVBVisitor& visitor,
     return KVBucket::Position(vbid);
 }
 
+size_t KVBucket::pauseResumeVisit(PauseResumeVBVisitor& visitor,
+                                  size_t currentPosition,
+                                  gsl::span<Vbid> vbsToVisit) {
+    for (; currentPosition < vbsToVisit.size(); ++currentPosition) {
+        VBucketPtr vb = getVBucket(vbsToVisit[currentPosition]);
+        if (vb) {
+            bool visitorPaused = !visitor.visit(*vb);
+            if (visitorPaused) {
+                break;
+            }
+        }
+    }
+    return currentPosition;
+}
+
 KVBucket::Position KVBucket::startPosition() const
 {
     return KVBucket::Position(Vbid(0));

engines/ep/src/kv_bucket.h

@@ -382,6 +382,24 @@ class KVBucket : public KVBucketIface {
                               Position& start_pos,
                               VBucketFilter* filter = nullptr) override;
 
+    /**
+     * Visit vbuckets in the order specified,
+     * allowing the visitor to indicate a pause.
+     *
+     * @param visitor object with a <code>virtual bool visit(VBucket& vb)</code>
+     *                method which returns true if the vb has been processed
+     *                and false if iteration should be paused
+     * @param currentPosition the vbsToVisit index at which iteration should
+     *                        proceed
+     * @param vbsToVisit span of vbids to visit in order
+     * @return the next position to visit
+     *         (the vbsToVisit index at which the visitor indicated a pause)
+     *         or <code>vbsToVisit.size()</code> at completion
+     */
+    [[nodiscard]] size_t pauseResumeVisit(PauseResumeVBVisitor& visitor,
+                                          size_t currentPosition,
+                                          gsl::span<Vbid> vbsToVisit);
+
     Position startPosition() const override;
 
     Position endPosition() const override;

engines/ep/src/warmup.cc

@@ -429,9 +429,8 @@ class WarmupBackfillTask : public GlobalTask {
                               120) /
                              100),
           currentNumBackfillTasks(threadTaskCount),
-          filter(warmup.shardVbIds[shardId]),
           visitor(bucket, *this),
-          epStorePosition(bucket.startPosition()) {
+          vbsToVisit(warmup.shardVbIds[shardId]) {
         warmup.addToTaskSet(uid);
     }
 
@@ -446,7 +445,7 @@ class WarmupBackfillTask : public GlobalTask {
     bool run() override {
         TRACE_EVENT1(
                 "ep-engine/task", "WarmupBackfillTask", "shard", getShardId());
-        if (filter.empty() || engine->getEpStats().isShutdown) {
+        if (vbsToVisit.empty() || engine->getEpStats().isShutdown) {
             // Technically "isShutdown" being true doesn't equate to a
             // successful task finish, however if we are shutting down we want
             // warmup to advance and be considered "done".
@@ -457,8 +456,7 @@ class WarmupBackfillTask : public GlobalTask {
         visitor.begin();
         auto& kvBucket = *engine->getKVBucket();
         try {
-            epStorePosition = kvBucket.pauseResumeVisit(
-                    visitor, epStorePosition, &filter);
+            visitPos = kvBucket.pauseResumeVisit(visitor, visitPos, vbsToVisit);
         } catch (std::exception& e) {
             EP_LOG_CRITICAL(
                     "WarmupBackfillTask::run(): caught exception while running "
@@ -467,7 +465,7 @@ class WarmupBackfillTask : public GlobalTask {
             finishTask(false);
             return false;
         }
-        if (epStorePosition == kvBucket.endPosition()) {
+        if (visitPos >= vbsToVisit.size()) {
             finishTask(true);
             return false;
         }
@@ -517,9 +515,9 @@ class WarmupBackfillTask : public GlobalTask {
     /// After how long should this task yield, allowing other tasks to run?
     const std::chrono::milliseconds maxExpectedRuntime;
     std::atomic<size_t>& currentNumBackfillTasks;
-    VBucketFilter filter;
     WarmupVbucketVisitor visitor;
-    KVBucketIface::Position epStorePosition;
+    size_t visitPos = 0;
+    std::vector<Vbid> vbsToVisit;
 };
 
 bool WarmupVbucketVisitor::visit(VBucket& vb) {

engines/ep/tests/module_tests/evp_store_warmup_test.cc

@@ -31,6 +31,7 @@
 #include "programs/engine_testapp/mock_cookie.h"
 #include "programs/engine_testapp/mock_server.h"
 #include "test_helpers.h"
+#include "vb_visitors.h"
 #include "vbucket.h"
 #include "vbucket_state.h"
 #include "warmup.h"
@@ -718,6 +719,64 @@ TEST_F(WarmupTest, MB_58135_CorruptAccessLog) {
     EXPECT_EQ(WarmupState::State::Done, warmup->getWarmupState());
 }
 
+class WarmupOrderTest : public SingleThreadedKVBucketTest {
+    void SetUp() override {
+        config_string = "max_vbuckets=32";
+        SingleThreadedKVBucketTest::SetUp();
+        for (size_t ii = 0; ii < store->getVBMapSize(); ++ii) {
+            setVBucketStateAndRunPersistTask(
+                    Vbid(ii),
+                    (ii & 4) ? vbucket_state_active : vbucket_state_replica);
+        }
+    }
+};
+
+/// Test that KVBucket::pauseResumeVisit visits vbuckets in the order specified
+TEST_F(WarmupOrderTest, pauseResumeVisit) {
+    class Visitor final : public PauseResumeVBVisitor {
+    public:
+        bool visit(VBucket& vb) override {
+            auto vbid = vb.getId();
+            if (pause && *pause == vbid) {
+                return false;
+            }
+            visited.push_back(vbid);
+            return true;
+        }
+
+        std::optional<Vbid> pause;
+        std::vector<Vbid> visited;
+    } visitor;
+
+    std::vector<Vbid> toVisit;
+    for (int ii = store->getVBMapSize() - 1; ii >= 0; ii -= 3) {
+        toVisit.emplace_back(ii);
+    }
+    EXPECT_EQ(11, toVisit.size());
+
+    // Visit all
+    size_t position = 0;
+    position = store->pauseResumeVisit(visitor, position, toVisit);
+    EXPECT_EQ(toVisit.size(), position);
+    EXPECT_EQ(toVisit, visitor.visited);
+
+    // Visit until the middle and pause
+    const auto pausePosition = toVisit.size() / 2;
+    visitor.pause = toVisit.at(pausePosition);
+    visitor.visited.clear();
+    position = 0;
+    position = store->pauseResumeVisit(visitor, position, toVisit);
+    EXPECT_EQ(pausePosition, position);
+    EXPECT_EQ(gsl::span(toVisit).subspan(0, pausePosition),
+              gsl::span(visitor.visited));
+
+    // Continue from middle
+    visitor.pause.reset();
+    position = store->pauseResumeVisit(visitor, position, toVisit);
+    EXPECT_EQ(position, toVisit.size());
+    EXPECT_EQ(visitor.visited, toVisit);
+}
+
 // Test fixture for Durability-related Warmup tests.
 class DurabilityWarmupTest : public DurabilityKVBucketTest {
 protected: