MB-69328: Add FBR/Fusion stat to interlock with cache transfer

Provide a new stat group and stat to return if the cache transfer
has completed and the vbucket move can continue.

vb_1:snapshot_rebalance_continue true

The expectation is that ns_server is only reading/using this stat
after a vbucket was created from a snapshot. E.g. when a setvbucket
used the "use_snapshot=fbr|fusion" option. In those cases the
vbucket is expected to have an empty cache and require cache
population (network cache transfer). The flag remains false until
KV concludes the transfer (CacheTransferEnd received). ns_server
will wait on this stat in the fusion/fbr vbucket move path before
proceeding to takover or before concluding a replica move is
ready.

If a fbr/fusion vbucket changes state we also flip this flag to
true.

For any vbucket created by warmup or setvbucket then this value
defaults to true.

Change-Id: I66354b4f44d49b543612989497613328ab8ccdfe
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/236359
Reviewed-by: Trond Norbye <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: jimwwalker

daemon/buckets_test.cc

@@ -24,9 +24,9 @@ TEST(BucketTest, Reset) {
     // check the main ones we run against.
     static constexpr size_t expectedBucketSize =
 #if defined(__linux) && defined(__x86_64__)
-            5680;
+            5696;
 #elif defined(__APPLE__)
-            5776;
+            5792;
 #else
             0;
 #endif

engines/ep/src/dcp/consumer.cc

@@ -365,6 +365,14 @@ cb::engine_errc DcpConsumer::addStream(uint32_t opaque,
                                snap_end_seqno,
                                high_seqno,
                                vb_manifest_uid);
+
+    // This stream never did a cache transfer we can flag any snapshot type
+    // rebalance can continue
+    if (!isFlagSet(stream->getFlags(),
+                   cb::mcbp::DcpAddStreamFlag::CacheTransfer)) {
+        vb->setSnapshotRebalanceCanContinue();
+    }
+
     registerStream(stream);
     readyStreamsVBQueue.lock()->push_back(vbucket);
     opaqueMap_[new_opaque] = std::make_pair(opaque, vbucket);

engines/ep/src/dcp/passive_stream.cc

@@ -1530,14 +1530,8 @@ cb::engine_errc PassiveStream::processCacheTransferEnd(
         return cb::engine_errc::not_my_vbucket;
     }
 
-    // Log this for now, will remove and switch to stats later.
-    OBJ_LOG_INFO_CTX(*this,
-                     "PassiveStream::processCacheTransferEnd: Cache transfer "
-                     "end received",
-                     {"vb", vb_});
+    vb->setSnapshotRebalanceCanContinue();
 
-    // @todo: update some state (Vbucket?) to indicate that the cache transfer
-    // is complete.
     return cb::engine_errc::success;
 }
 

engines/ep/src/ep_bucket.cc

@@ -3312,6 +3312,42 @@ cb::engine_errc EPBucket::doSnapshotStatus(const StatCollector& collector,
     return cb::engine_errc::success;
 }
 
+cb::engine_errc EPBucket::doSnapshotMoveStats(const StatCollector& collector,
+                                              std::string_view input) {
+    const std::string_view stat_key = "snapshot-move ";
+    // Configure for 1 vbid or all known vbuckets
+    std::vector<Vbid> ids;
+    if (input.size() > stat_key.size()) {
+        // input should be a vbid
+        input.remove_prefix(stat_key.size());
+        uint16_t vbucket_id(0);
+        if (!safe_strtous(input, vbucket_id)) {
+            return cb::engine_errc::invalid_arguments;
+        }
+        ids.emplace_back(vbucket_id);
+    } else {
+        // Populate for all possible vbuckets.
+        for (size_t i = 0; i < vbMap.getSize(); i++) {
+            ids.emplace_back(i);
+        }
+    }
+
+    for (const auto id : ids) {
+        auto vb = getVBucket(id);
+        if (!vb) {
+            continue;
+        }
+        fmt::memory_buffer key;
+        fmt::format_to(std::back_inserter(key),
+                       "vb_{}:snapshot_rebalance_continue",
+                       id.get());
+        collector.addStat(std::string_view(key.data(), key.size()),
+                          vb->canSnapshotRebalanceContinue());
+    }
+
+    return cb::engine_errc::success;
+}
+
 cb::engine_errc EPBucket::doSnapshotDeks(const StatCollector& collector) {
     // Populate for all possible vbuckets.
     for (Vbid::id_type i = 0; i < vbMap.getSize(); i++) {

engines/ep/src/ep_bucket.h

@@ -490,6 +490,13 @@ class EPBucket : public KVBucket {
     cb::engine_errc doSnapshotStatus(const StatCollector&,
                                      std::string_view) override;
 
+    /**
+     * Handle the brief snapshot-move stat which is used by ns_server to
+     * monitor moves.
+     */
+    cb::engine_errc doSnapshotMoveStats(const StatCollector&,
+                                        std::string_view) override;
+
     /**
      * Handle the brief snapshot-deks stat which is used by ns_server to
      * determine which deks are in use in snapshots. Deks are returned

engines/ep/src/ep_engine.cc

@@ -5340,6 +5340,9 @@ cb::engine_errc EventuallyPersistentEngine::getStats(
     if (key == "snapshot-deks"sv) {
         return getKVBucket()->doSnapshotDeks(bucketCollector);
     }
+    if (key.starts_with("snapshot-move")) {
+        return getKVBucket()->doSnapshotMoveStats(bucketCollector, key);
+    }
 
     // Unknown stat requested
     return cb::engine_errc::no_such_key;

engines/ep/src/ep_vb.cc

@@ -47,6 +47,19 @@
 #include <statistics/collector.h>
 #include <utilities/logtags.h>
 
+static bool getInitialSnapshotRebalanceCanContinue(
+        CreateVbucketMethod creationMethod) {
+    switch (creationMethod) {
+    case CreateVbucketMethod::FBR:
+    case CreateVbucketMethod::Fusion:
+        return false;
+    case CreateVbucketMethod::Warmup:
+    case CreateVbucketMethod::SetVbucket:
+        return true;
+    }
+    folly::assume_unreachable();
+}
+
 EPVBucket::EPVBucket(Vbid i,
                      vbucket_state_t newState,
                      EPStats& st,
@@ -103,7 +116,9 @@ EPVBucket::EPVBucket(Vbid i,
               maxPrepareSeqno),
       shard(kvshard),
       rangeScans(static_cast<EPBucket*>(bucket), *this),
-      canReceiveCacheTransfer(newState == vbucket_state_replica) {
+      canReceiveCacheTransfer(newState == vbucket_state_replica),
+      snapshotRebalanceCanContinue(
+              getInitialSnapshotRebalanceCanContinue(creationMethod)) {
     if (config.isBfilterEnabled()) {
         bFilterData.lock()->kvStoreBfilterEnabled = true;
     }

engines/ep/src/ep_vb.h

@@ -387,6 +387,14 @@ class EPVBucket : public VBucket {
         canReceiveCacheTransfer = false;
     }
 
+    void setSnapshotRebalanceCanContinue() override {
+        snapshotRebalanceCanContinue = true;
+    }
+
+    bool canSnapshotRebalanceContinue() const override {
+        return snapshotRebalanceCanContinue.load();
+    }
+
 protected:
     /**
      * queue a background fetch of the specified item.
@@ -564,6 +572,10 @@ class EPVBucket : public VBucket {
     /// Flag to indicate if the vbucket can receive a DCP cache transfer.
     std::atomic<bool> canReceiveCacheTransfer{true};
 
+    /// Flag to indicate if the vbucket can proceed in a snapshot based
+    /// rebalance, i.e. FBR/Fusion
+    std::atomic<bool> snapshotRebalanceCanContinue{false};
+
     friend class EPVBucketTest;
 };
 

engines/ep/src/kv_bucket_iface.h

@@ -985,6 +985,11 @@ class KVBucketIface {
         return cb::engine_errc::not_supported;
     }
 
+    [[nodiscard]] virtual cb::engine_errc doSnapshotMoveStats(
+            const StatCollector&, std::string_view) {
+        return cb::engine_errc::not_supported;
+    }
+
     [[nodiscard]] virtual cb::engine_errc doSnapshotDeks(const StatCollector&) {
         return cb::engine_errc::not_supported;
     }

engines/ep/src/vbucket.cc

@@ -699,7 +699,6 @@ void VBucket::setState_UNLOCKED(
         const nlohmann::json* meta,
         const std::unique_lock<folly::SharedMutex>& vbStateLock) {
     vbucket_state_t oldstate = state;
-
     const bool changingState = to != oldstate;
 
     // Validate (optional) meta content.
@@ -747,6 +746,10 @@ void VBucket::setState_UNLOCKED(
         // to replica and do a cache transfer in situations we've not
         // considered.
         disableCacheTransfer();
+
+        // Any state change and this vbucket should no longer block a future
+        // rebalance
+        setSnapshotRebalanceCanContinue();
     }
 }
 
@@ -3471,6 +3474,10 @@ void VBucket::_addStats(VBucketStatsDetailLevel detail,
         addStat("persistence_seqno", getPersistenceSeqno(), add_stat, c);
         hlc.addStats(statPrefix, add_stat, c);
         addStat("creation_method", to_string(creationMethod), add_stat, c);
+        addStat("snapshot_rebalance_continue",
+                canSnapshotRebalanceContinue(),
+                add_stat,
+                c);
     }
         // fallthrough
     case VBucketStatsDetailLevel::Durability: