[BP] MB-59817 [2/2]: Warmup active vbuckets before other

...to prioritise DGM workload efficiency over faster failover.

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

Author: pavlosg

engines/ep/src/warmup.cc

@@ -34,11 +34,9 @@
 #include <platform/timeutils.h>
 #include <statistics/cbstat_collector.h>
 #include <utilities/logtags.h>
-#include <array>
 #include <limits>
 #include <memory>
 #include <optional>
-#include <random>
 #include <string>
 #include <utility>
 
@@ -2047,70 +2045,30 @@ uint16_t Warmup::getNumKVStores() {
 }
 
 void Warmup::populateShardVbStates() {
-    uint16_t numKvs = getNumKVStores();
-
-    for (size_t i = 0; i < numKvs; i++) {
-        std::vector<vbucket_state*> kvStoreVbStates =
-                store.getRWUnderlyingByShard(i)->listPersistedVbuckets();
-        for (uint16_t j = 0; j < kvStoreVbStates.size(); j++) {
-            if (!kvStoreVbStates[j]) {
+    const uint16_t numShards = getNumKVStores();
+    for (uint16_t shardIdx = 0; shardIdx < numShards; ++shardIdx) {
+        const std::vector<vbucket_state*> curShardStates =
+                store.getRWUnderlyingByShard(shardIdx)->listPersistedVbuckets();
+        auto& statesMap = shardVbStates[shardIdx];
+        auto& statesVec = shardVbIds[shardIdx];
+
+        for (uint16_t vbIdx = 0; vbIdx < curShardStates.size(); ++vbIdx) {
+            if (!curShardStates[vbIdx]) {
                 continue;
             }
-            auto vb = (j * numKvs) + i;
-            std::map<Vbid, vbucket_state>& shardVB =
-                    shardVbStates[vb % store.vbMap.getNumShards()];
-            shardVB.insert(std::pair<Vbid, vbucket_state>(
-                    Vbid(vb), *(kvStoreVbStates[j])));
+            const Vbid vbid(vbIdx * numShards + shardIdx);
+            statesMap.insert({vbid, *curShardStates[vbIdx]});
         }
-    }
 
-    for (size_t i = 0; i < store.vbMap.shards.size(); i++) {
-        std::vector<Vbid> activeVBs, otherVBs;
-        std::map<Vbid, vbucket_state>::const_iterator it;
-        for (auto shardIt : shardVbStates[i]) {
-            Vbid vbid = shardIt.first;
-            vbucket_state vbs = shardIt.second;
-            if (vbs.transition.state == vbucket_state_active) {
-                activeVBs.push_back(vbid);
-            } else {
-                otherVBs.push_back(vbid);
+        // First push all active vbuckets and then the rest
+        for (const auto& item : statesMap) {
+            if (item.second.transition.state == vbucket_state_active) {
+                statesVec.push_back(item.first);
             }
         }
-
-        // Push one active VB to the front.
-        // When the ratio of RAM to VBucket is poor (big vbuckets) this will
-        // ensure we at least bring active data in before replicas eat RAM.
-        if (!activeVBs.empty()) {
-            shardVbIds[i].push_back(activeVBs.back());
-            activeVBs.pop_back();
-        }
-
-        // Now the VB lottery can begin.
-        // Generate a psudeo random, weighted list of active/replica vbuckets.
-        // The random seed is the shard ID so that re-running warmup
-        // for the same shard and vbucket set always gives the same output and keeps
-        // nodes of the cluster more equal after a warmup.
-
-        std::mt19937 twister(i);
-        // Give 'true' (aka active) 60% of the time
-        // Give 'false' (aka other) 40% of the time.
-        std::bernoulli_distribution distribute(0.6);
-        std::array<std::vector<Vbid>*, 2> activeReplicaSource = {
-                {&activeVBs, &otherVBs}};
-
-        while (!activeVBs.empty() || !otherVBs.empty()) {
-            const bool active = distribute(twister);
-            int num = active ? 0 : 1;
-            if (!activeReplicaSource[num]->empty()) {
-                shardVbIds[i].push_back(activeReplicaSource[num]->back());
-                activeReplicaSource[num]->pop_back();
-            } else {
-                // Once active or replica set is empty, just drain the other one.
-                num = num ^ 1;
-                while (!activeReplicaSource[num]->empty()) {
-                    shardVbIds[i].push_back(activeReplicaSource[num]->back());
-                    activeReplicaSource[num]->pop_back();
-                }
+        for (const auto& item : statesMap) {
+            if (item.second.transition.state != vbucket_state_active) {
+                statesVec.push_back(item.first);
             }
         }
     }

engines/ep/src/warmup.h

@@ -323,6 +323,10 @@ class Warmup {
         return "running";
     }
 
+    const std::vector<std::vector<Vbid>>& getShardVbIds() const {
+        return shardVbIds;
+    }
+
     /**
      * Testing hook which if set is called every time warmup transitions to
      * a new state.

engines/ep/tests/module_tests/evp_store_warmup_test.cc

@@ -777,6 +777,29 @@ TEST_F(WarmupOrderTest, pauseResumeVisit) {
     EXPECT_EQ(visitor.visited, toVisit);
 }
 
+/// Test that all the active vbuckets in a shard are loaded first
+TEST_F(WarmupOrderTest, populateShardVbStates) {
+    resetEngineAndEnableWarmup();
+    runReadersUntilWarmedUp();
+    const auto* warmup = store->getWarmup();
+    ASSERT_TRUE(warmup);
+    for (const auto& shard : warmup->getShardVbIds()) {
+        EXPECT_EQ(16, shard.size());
+        bool seenNotActive = false;
+        for (const auto& vbid : shard) {
+            const auto vb = store->getVBucket(vbid);
+            ASSERT_TRUE(vb);
+            const auto state = vb->getState();
+            if (seenNotActive) {
+                EXPECT_NE(vbucket_state_active, state);
+            }
+            if (state != vbucket_state_active) {
+                seenNotActive = true;
+            }
+        }
+    }
+}
+
 // Test fixture for Durability-related Warmup tests.
 class DurabilityWarmupTest : public DurabilityKVBucketTest {
 protected: