MB-41804: Track correct number of vbuckets in each state

When moving between states, the vbmap needs to decrement the count of
vbuckets in the old state, and increment in the new state. It
previously only decremented the old state.

Change-Id: I2f2f50e87cbd0c09f8496cfddeb4e7726339d4da
Reviewed-on: http://review.couchbase.org/c/kv_engine/+/139835
Tested-by: Build Bot <[email protected]>
Well-Formed: Build Bot <[email protected]>
Reviewed-by: Ben Huddleston <[email protected]>

Author: jameseh96

engines/ep/src/kv_bucket.cc

@@ -802,6 +802,7 @@ ENGINE_ERROR_CODE KVBucket::setVBucketState_UNLOCKED(
         } else {
             vb->setState(to);
         }
+        vbMap.incrVBStateCount(to);
 
         if (oldstate != to && notify_dcp) {
             bool closeInboundStreams = false;

engines/ep/src/vbucketmap.h

@@ -109,6 +109,14 @@ friend class Warmup;
     void setHLCDriftAheadThreshold(std::chrono::microseconds threshold);
     void setHLCDriftBehindThreshold(std::chrono::microseconds threshold);
 
+    /**
+     * Increment the vb count for the given state.
+     * @param state the state for which the vb count is to be incremented.
+     */
+    void incrVBStateCount(vbucket_state_t state) {
+        ++vbStateCount[state];
+    }
+
     /**
      * Decrement the vb count for the given state.
      * @param state  the state for which the vb count is to be decremented.

engines/ep/tests/module_tests/kv_bucket_test.cc

@@ -1353,6 +1353,31 @@ TEST_P(KVBucketParamTest, MB_34346) {
             << "Should still have 0 items after time-travelling/expiry";
 }
 
+TEST_P(KVBucketParamTest, VbucketStateCounts) {
+    // confirm the vbMap correctly changes the number of vbuckets in a given
+    // state when vbuckets change state
+    auto vbA = 0;
+    auto vbB = 1;
+
+    auto expectVbCounts = [this](uint16_t active, uint16_t replica) {
+        auto message = "Expected " + std::to_string(active) + " active and " +
+                       std::to_string(replica) + " replica vbs";
+        EXPECT_EQ(active, store->getNumOfVBucketsInState(vbucket_state_active))
+                << message;
+        EXPECT_EQ(replica,
+                  store->getNumOfVBucketsInState(vbucket_state_replica))
+                << message;
+    };
+    store->setVBucketState(vbA, vbucket_state_active, false);
+    expectVbCounts(1, 0);
+    store->setVBucketState(vbB, vbucket_state_active, false);
+    expectVbCounts(2, 0);
+    store->setVBucketState(vbA, vbucket_state_replica, false);
+    expectVbCounts(1, 1);
+    store->setVBucketState(vbB, vbucket_state_replica, false);
+    expectVbCounts(0, 2);
+}
+
 class StoreIfTest : public KVBucketTest {
 public:
     void SetUp() override {