MB-48038: Implement CheckpointVisitor::getVBucketComparator

In the follow-up patch I will move the checkpoint expel/removal logic
from the CheckpointRemoverTask to the CheckpointVisitor.

In preparation for that, allow the Visitor to visit vbucket in "highest
checkpoint mem-usage" order. Which is the same as what the RemoverTask
currently does.

Change-Id: I8e9e3dd007f2344122a5fb089328ada5a36a7619
Reviewed-on: http://review.couchbase.org/c/kv_engine/+/160317
Tested-by: Build Bot <build@couchbase.com>
Reviewed-by: Dave Rigby <daver@couchbase.com>

engines/ep/src/checkpoint_visitor.cc

@@ -66,4 +66,26 @@ CheckpointVisitor::shouldInterrupt() {
 
     // Rely on the default behaviour otherwise
     return CappedDurationVBucketVisitor::shouldInterrupt();
+}
+
+std::function<bool(const Vbid&, const Vbid&)>
+CheckpointVisitor::getVBucketComparator() const {
+    // Some ep_testsuite failures highlight that accessing vbucket in the VBMap
+    // within the comparator may cause issues as the VBMap may change while
+    // the comparator is being called. Thus, we build-up a vb-ckpt-mem-usage
+    // vector from the current state of VBMap and then we pass it and use it
+    // in the comparator.
+
+    const auto& vbMap = store->getVBuckets();
+    std::vector<size_t> ckptMemUsage(vbMap.getSize());
+    const auto vbuckets = vbMap.getBuckets();
+    for (const auto vbid : vbuckets) {
+        const auto vb = store->getVBucket(vbid);
+        ckptMemUsage[vbid.get()] = vb ? vb->getChkMgrMemUsage() : 0;
+    }
+
+    return [ckptMemUsage = std::move(ckptMemUsage)](const Vbid& vbid1,
+                                                    const Vbid& vbid2) -> bool {
+        return ckptMemUsage.at(vbid1.get()) < ckptMemUsage.at(vbid2.get());
+    };
 }

engines/ep/src/checkpoint_visitor.h

@@ -39,6 +39,14 @@ class CheckpointVisitor : public CappedDurationVBucketVisitor {
      */
     ExecutionState shouldInterrupt() override;
 
+    /**
+     * @return The comparator function that the Adaptor has to use for ordering
+     *  vbuckets to visit. The CheckpointVisitor provides a comparator that
+     *  orders from the highest to the lowest CM memory usage.
+     */
+    std::function<bool(const Vbid&, const Vbid&)> getVBucketComparator()
+            const override;
+
 private:
     KVBucketIface* store;
     EPStats& stats;

engines/ep/src/vbucketmap.cc

@@ -72,20 +72,6 @@ std::vector<Vbid> VBucketMap::getBuckets() const {
     return rv;
 }
 
-std::vector<Vbid> VBucketMap::getBucketsSortedByState() const {
-    std::vector<Vbid> rv;
-    for (int state = vbucket_state_active;
-         state <= vbucket_state_dead; ++state) {
-        for (size_t i = 0; i < size; ++i) {
-            VBucketPtr b = getBucket(Vbid(i));
-            if (b && b->getState() == state) {
-                rv.push_back(b->getId());
-            }
-        }
-    }
-    return rv;
-}
-
 std::vector<Vbid> VBucketMap::getBucketsInState(vbucket_state_t state) const {
     std::vector<Vbid> rv;
     for (size_t i = 0; i < size; ++i) {

engines/ep/src/vbucketmap.h

@@ -75,7 +75,6 @@ friend class Warmup;
         return size;
     }
     std::vector<Vbid> getBuckets() const;
-    std::vector<Vbid> getBucketsSortedByState() const;
 
     /**
      * Returns a vector containing the vbuckets from the map that are in the

engines/ep/tests/module_tests/checkpoint_remover_test.cc

@@ -19,6 +19,7 @@
 #include "checkpoint_manager.h"
 #include "checkpoint_remover.h"
 #include "checkpoint_utils.h"
+#include "checkpoint_visitor.h"
 #include "collections/vbucket_manifest_handles.h"
 #include "dcp/response.h"
 #include "test_helpers.h"
@@ -59,6 +60,39 @@ TEST_F(CheckpointRemoverEPTest, GetVBucketsSortedByChkMgrMem) {
     }
 }
 
+/**
+ * Check that CheckpointVisitor orders vbuckets to visit by "highest checkpoint
+ * mem-usage" order.
+ */
+TEST_F(CheckpointRemoverEPTest, CheckpointVisitorVBucketOrder) {
+    std::deque<Vbid> vbuckets;
+    for (uint16_t i = 0; i < 3; i++) {
+        setVBucketStateAndRunPersistTask(Vbid(i), vbucket_state_active);
+        for (uint16_t j = 0; j < i; j++) {
+            std::string doc_key =
+                    "key_" + std::to_string(i) + "_" + std::to_string(j);
+            store_item(Vbid(i), makeStoredDocKey(doc_key), "value");
+        }
+        vbuckets.emplace_back(vbid);
+    }
+
+    std::atomic<bool> stateFinalizer = true;
+    auto visitor = CheckpointVisitor(
+            store, engine->getEpStats(), stateFinalizer, 1 /*memToClear*/);
+    auto comparator = visitor.getVBucketComparator();
+    std::sort(vbuckets.begin(), vbuckets.end(), comparator);
+
+    ASSERT_EQ(3, vbuckets.size());
+    for (size_t i = 1; i < vbuckets.size(); i++) {
+        auto thisVbid = vbuckets[i];
+        auto prevVbid = vbuckets[i - 1];
+        // This vBucket should have a greater memory usage than the one previous
+        // to it in the map
+        ASSERT_GE(store->getVBucket(thisVbid)->getChkMgrMemUsage(),
+                  store->getVBucket(prevVbid)->getChkMgrMemUsage());
+    }
+}
+
 #ifndef WIN32
 /**
  * Check that the CheckpointManager memory usage calculation is correct and