SERVER-66698 Clear BucketCatalog namespaces lazily using clear registry

4 files changed, 103 insertions, 49 deletions

diff --git a/src/mongo/db/op_observer/op_observer_impl.cpp b/src/mongo/db/op_observer/op_observer_impl.cpp
index d065b89e932..8c498e58d69 100644
--- a/src/mongo/db/op_observer/op_observer_impl.cpp
+++ b/src/mongo/db/op_observer/op_observer_impl.cpp
@@ -2400,9 +2400,10 @@ void OpObserverImpl::_onReplicationRollback(OperationContext* opCtx,
             timeseriesNamespaces.insert(ns.getTimeseriesViewNamespace());
         }
     }
-    BucketCatalog::get(opCtx).clear([&timeseriesNamespaces](const NamespaceString& bucketNs) {
-        return timeseriesNamespaces.contains(bucketNs);
-    });
+    BucketCatalog::get(opCtx).clear(
+        [timeseriesNamespaces = std::move(timeseriesNamespaces)](const NamespaceString& bucketNs) {
+            return timeseriesNamespaces.contains(bucketNs);
+        });
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/timeseries/bucket_catalog.cpp b/src/mongo/db/timeseries/bucket_catalog.cpp
index e11d5b855ec..6af3ae04f95 100644
--- a/src/mongo/db/timeseries/bucket_catalog.cpp
+++ b/src/mongo/db/timeseries/bucket_catalog.cpp
@@ -373,8 +373,7 @@ void BucketCatalog::EraManager::_incrementEraCountHelper(uint64_t era) {
     }
 }
 
-bool BucketCatalog::EraManager::hasBeenCleared(Bucket* bucket) {
-    stdx::lock_guard lk{*_mutex};
+bool BucketCatalog::EraManager::hasBeenCleared(WithLock catalogLock, Bucket* bucket) {
     for (auto it = _clearRegistry.find(bucket->getEra() + 1); it != _clearRegistry.end(); ++it) {
         if (it->second(bucket->_ns)) {
             return true;
@@ -703,7 +702,7 @@ Status BucketCatalog::reopenBucket(OperationContext* opCtx,
     return Status::OK();
 }
 
-BSONObj BucketCatalog::getMetadata(const BucketHandle& handle) const {
+BSONObj BucketCatalog::getMetadata(const BucketHandle& handle) {
     auto const& stripe = _stripes[handle.stripe];
     stdx::lock_guard stripeLock{stripe.mutex};
 
@@ -847,7 +846,11 @@ void BucketCatalog::abort(std::shared_ptr<WriteBatch> batch, const Status& statu
 }
 
 void BucketCatalog::clear(const OID& oid) {
-    auto result = _setBucketState(oid, BucketState::kCleared);
+    boost::optional<BucketState> result;
+    {
+        stdx::lock_guard catalogLock{_mutex};
+        result = _setBucketState(catalogLock, oid, BucketState::kCleared);
+    }
     if (result && *result == BucketState::kPreparedAndCleared) {
         hangTimeseriesDirectModificationBeforeWriteConflict.pauseWhileSet();
         throwWriteConflictException("Prepared bucket can no longer be inserted into.");
@@ -855,6 +858,12 @@ void BucketCatalog::clear(const OID& oid) {
 }
 
 void BucketCatalog::clear(ShouldClearFn&& shouldClear) {
+    if (feature_flags::gTimeseriesScalabilityImprovements.isEnabled(
+            serverGlobalParams.featureCompatibility)) {
+        uint64_t era = _eraManager.incrementEra();
+        _eraManager.insertToRegistry(era, std::move(shouldClear));
+        return;
+    }
     for (auto& stripe : _stripes) {
         stdx::lock_guard stripeLock{stripe.mutex};
         for (auto it = stripe.allBuckets.begin(); it != stripe.allBuckets.end();) {
@@ -876,20 +885,16 @@ void BucketCatalog::clear(ShouldClearFn&& shouldClear) {
             it = nextIt;
         }
     }
-
-    uint64_t era = _eraManager.incrementEra();
-    if (feature_flags::gTimeseriesScalabilityImprovements.isEnabled(
-            serverGlobalParams.featureCompatibility)) {
-        _eraManager.insertToRegistry(era, std::move(shouldClear));
-    }
 }
 
 void BucketCatalog::clear(const NamespaceString& ns) {
-    clear([&ns](const NamespaceString& bucketNs) { return bucketNs == ns; });
+    clear([ns](const NamespaceString& bucketNs) { return bucketNs == ns; });
 }
 
 void BucketCatalog::clear(StringData dbName) {
-    clear([&dbName](const NamespaceString& bucketNs) { return bucketNs.db() == dbName; });
+    clear([dbName = dbName.toString()](const NamespaceString& bucketNs) {
+        return bucketNs.db() == dbName;
+    });
 }
 
 void BucketCatalog::_appendExecutionStatsToBuilder(const ExecutionStats* stats,
@@ -1014,14 +1019,14 @@ BucketCatalog::StripeNumber BucketCatalog::_getStripeNumber(const BucketKey& key
 const BucketCatalog::Bucket* BucketCatalog::_findBucket(const Stripe& stripe,
                                                         WithLock,
                                                         const OID& id,
-                                                        ReturnClearedBuckets mode) const {
+                                                        ReturnClearedBuckets mode) {
     auto it = stripe.allBuckets.find(id);
     if (it != stripe.allBuckets.end()) {
         if (mode == ReturnClearedBuckets::kYes) {
             return it->second.get();
         }
 
-        auto state = _getBucketState(id);
+        auto state = _getBucketState(it->second.get());
         if (state && state != BucketState::kCleared && state != BucketState::kPreparedAndCleared) {
             return it->second.get();
         }
@@ -1042,7 +1047,7 @@ BucketCatalog::Bucket* BucketCatalog::_useBucketInState(Stripe* stripe,
                                                         BucketState targetState) {
     auto it = stripe->allBuckets.find(id);
     if (it != stripe->allBuckets.end()) {
-        auto state = _setBucketState(it->second->_id, targetState);
+        auto state = _setBucketState(it->second.get(), targetState);
         if (state && state != BucketState::kCleared && state != BucketState::kPreparedAndCleared) {
             return it->second.get();
         }
@@ -1063,7 +1068,7 @@ BucketCatalog::Bucket* BucketCatalog::_useBucket(Stripe* stripe,
 
     Bucket* bucket = it->second;
 
-    auto state = _getBucketState(bucket->id());
+    auto state = _getBucketState(bucket);
     if (state == BucketState::kNormal || state == BucketState::kPrepared) {
         _markBucketNotIdle(stripe, stripeLock, bucket);
         return bucket;
@@ -1725,15 +1730,30 @@ void BucketCatalog::_eraseBucketState(const OID& id) {
     _bucketStates.erase(id);
 }
 
-boost::optional<BucketCatalog::BucketState> BucketCatalog::_getBucketState(const OID& id) const {
+boost::optional<BucketCatalog::BucketState> BucketCatalog::_getBucketState(Bucket* bucket) {
     stdx::lock_guard catalogLock{_mutex};
-    auto it = _bucketStates.find(id);
+    // If the bucket has been cleared, we will set the bucket state accordingly to reflect that
+    // (kPreparedAndCleared or kCleared).
+    if (_eraManager.hasBeenCleared(catalogLock, bucket)) {
+        return _setBucketState(catalogLock, bucket->id(), BucketState::kCleared);
+    }
+    auto it = _bucketStates.find(bucket->id());
     return it != _bucketStates.end() ? boost::make_optional(it->second) : boost::none;
 }
 
-boost::optional<BucketCatalog::BucketState> BucketCatalog::_setBucketState(const OID& id,
+boost::optional<BucketCatalog::BucketState> BucketCatalog::_setBucketState(Bucket* bucket,
                                                                            BucketState target) {
     stdx::lock_guard catalogLock{_mutex};
+    if (_eraManager.hasBeenCleared(catalogLock, bucket)) {
+        return _setBucketState(catalogLock, bucket->id(), BucketState::kCleared);
+    }
+
+    return _setBucketState(catalogLock, bucket->id(), target);
+}
+
+boost::optional<BucketCatalog::BucketState> BucketCatalog::_setBucketState(WithLock catalogLock,
+                                                                           const OID& id,
+                                                                           BucketState target) {
     auto it = _bucketStates.find(id);
     if (it == _bucketStates.end()) {
         return boost::none;
diff --git a/src/mongo/db/timeseries/bucket_catalog.h b/src/mongo/db/timeseries/bucket_catalog.h
index 8c72c8ac676..8705ceaa333 100644
--- a/src/mongo/db/timeseries/bucket_catalog.h
+++ b/src/mongo/db/timeseries/bucket_catalog.h
@@ -279,7 +279,7 @@ public:
      * Returns an empty document if the given bucket cannot be found or if this time-series
      * collection was not created with a metadata field name.
      */
-    BSONObj getMetadata(const BucketHandle& bucket) const;
+    BSONObj getMetadata(const BucketHandle& bucket);
 
     /**
      * Tries to insert 'doc' into a suitable bucket. If an open bucket is full (or has incompatible
@@ -548,7 +548,7 @@ protected:
         // Returns whether the Bucket has been marked as cleared by checking against the
         // clearRegistry. Advances Bucket's era up to current global era if the bucket has not been
         // cleared.
-        bool hasBeenCleared(Bucket* bucket);
+        bool hasBeenCleared(WithLock catalogLock, Bucket* bucket);
 
     protected:
         void _decrementEraCountHelper(uint64_t era);
@@ -760,7 +760,7 @@ protected:
     const Bucket* _findBucket(const Stripe& stripe,
                               WithLock stripeLock,
                               const OID& id,
-                              ReturnClearedBuckets mode = ReturnClearedBuckets::kNo) const;
+                              ReturnClearedBuckets mode = ReturnClearedBuckets::kNo);
 
     /**
      * Retrieve a bucket for write use.
@@ -941,9 +941,10 @@ protected:
     void _appendExecutionStatsToBuilder(const ExecutionStats* stats, BSONObjBuilder* builder) const;
 
     /**
-     * Retreives the bucket state if it is tracked in the catalog.
+     * Retrieves the bucket state if it is tracked in the catalog. Modifies the bucket state if the
+     * bucket is found to have been cleared.
      */
-    boost::optional<BucketState> _getBucketState(const OID& id) const;
+    boost::optional<BucketState> _getBucketState(Bucket* bucket);
 
     /**
      * Initializes state for the given bucket to kNormal.
@@ -956,6 +957,14 @@ protected:
     void _eraseBucketState(const OID& id);
 
     /**
+     * Checks whether the bucket has been cleared before changing the bucket state to the target
+     * state. If the bucket has been cleared, it will set the state to kCleared instead and ignore
+     * the target state. The return value, if set, is the final state of the bucket
+     * with the given id.
+     */
+    boost::optional<BucketState> _setBucketState(Bucket* bucket, BucketState target);
+
+    /**
      * Changes the bucket state, taking into account the current state, the specified target state,
      * and allowed state transitions. The return value, if set, is the final state of the bucket
      * with the given id; if no such bucket exists, the return value will not be set.
@@ -963,7 +972,9 @@ protected:
      * Ex. For a bucket with state kPrepared, and a target of kCleared, the return will be
      * kPreparedAndCleared.
      */
-    boost::optional<BucketState> _setBucketState(const OID& id, BucketState target);
+    boost::optional<BucketState> _setBucketState(WithLock withLock,
+                                                 const OID& id,
+                                                 BucketState target);
 
     /**
      * Calculates the marginal memory usage for an archived bucket. The
diff --git a/src/mongo/db/timeseries/bucket_catalog_era_test.cpp b/src/mongo/db/timeseries/bucket_catalog_era_test.cpp
index 8069b3a38c3..dbfae608d34 100644
--- a/src/mongo/db/timeseries/bucket_catalog_era_test.cpp
+++ b/src/mongo/db/timeseries/bucket_catalog_era_test.cpp
@@ -41,7 +41,7 @@ public:
     Bucket* createBucket(const CreationInfo& info) {
         auto ptr = _allocateBucket(&_stripes[info.stripe], withLock, info);
         ptr->setNamespace(info.key.ns);
-        ASSERT_FALSE(_eraManager.hasBeenCleared(ptr));
+        ASSERT_FALSE(_eraManager.hasBeenCleared(withLock, ptr));
         return ptr;
     }
 
@@ -58,7 +58,7 @@ public:
     }
 
     bool cannotAccessBucket(Bucket* bucket) {
-        if (_eraManager.hasBeenCleared(bucket)) {
+        if (_eraManager.hasBeenCleared(withLock, bucket)) {
             _removeBucket(&_stripes[bucket->stripe()], withLock, bucket, false);
             return true;
         } else {
@@ -66,7 +66,20 @@ public:
         }
     }
 
-    Stripe stripe;
+    void checkAndRemoveClearedBucket(Bucket* bucket, BucketKey bucketKey, WithLock withLock) {
+        auto a = _findBucket(_stripes[_getStripeNumber(bucketKey)],
+                             withLock,
+                             bucket->id(),
+                             ReturnClearedBuckets::kYes);
+        ASSERT(a == bucket);
+        auto b = _findBucket(_stripes[_getStripeNumber(bucketKey)],
+                             withLock,
+                             bucket->id(),
+                             ReturnClearedBuckets::kNo);
+        ASSERT(b == nullptr);
+        _removeBucket(&_stripes[_getStripeNumber(bucketKey)], withLock, bucket, false);
+    }
+
     WithLock withLock = WithLock::withoutLock();
     NamespaceString ns1{"db.test1"};
     NamespaceString ns2{"db.test2"};
@@ -104,8 +117,7 @@ TEST_F(BucketCatalogEraTest, EraAdvancesAsExpected) {
     // bucket era values should remain unchanged.
     clear(ns1);
     ASSERT_EQ(_eraManager.getEra(), 1);
-    // TODO (SERVER-66698): Add checks on the buckets' era values.
-    // ASSERT_EQ(b1->era(), 0);
+    ASSERT_EQ(bucket1->getEra(), 0);
 
     // When clearing buckets of one namespace, we expect the era of buckets of any other namespace
     // to not change.
@@ -117,21 +129,20 @@ TEST_F(BucketCatalogEraTest, EraAdvancesAsExpected) {
     clear(ns1);
     ASSERT_EQ(_eraManager.getEra(), 2);
     ASSERT_EQ(bucket3->getEra(), 1);
-    // TODO (SERVER-66698): Add checks on the buckets' era values.
-    // ASSERT_EQ(b1->era(), 0);
-    // ASSERT_EQ(b2->era(), 1);
+    ASSERT_EQ(bucket1->getEra(), 0);
+    ASSERT_EQ(bucket2->getEra(), 1);
 }
 
 TEST_F(BucketCatalogEraTest, EraCountMapUpdatedCorrectly) {
     RAIIServerParameterControllerForTest controller{"featureFlagTimeseriesScalabilityImprovements",
                                                     true};
 
-    // TODO (SERVER-66698): Change count assertions now that Buckets are cleared lazily.
     // Creating a bucket in a new era should add a counter for that era to the map.
     auto bucket1 = createBucket(info1);
     ASSERT_EQ(bucket1->getEra(), 0);
     ASSERT_EQ(_eraManager.getCountForEra(0), 1);
     clear(ns1);
+    checkAndRemoveClearedBucket(bucket1, bucketKey1, withLock);
 
     // When the last bucket in an era is destructed, the counter in the map should be removed.
     ASSERT_EQ(_eraManager.getCountForEra(0), 0);
@@ -144,6 +155,7 @@ TEST_F(BucketCatalogEraTest, EraCountMapUpdatedCorrectly) {
     ASSERT_EQ(bucket3->getEra(), 1);
     ASSERT_EQ(_eraManager.getCountForEra(1), 2);
     clear(ns2);
+    checkAndRemoveClearedBucket(bucket3, bucketKey2, withLock);
     ASSERT_EQ(_eraManager.getCountForEra(1), 1);
 
     // A bucket in one era being destroyed and the counter decrementing should not affect a
@@ -152,6 +164,7 @@ TEST_F(BucketCatalogEraTest, EraCountMapUpdatedCorrectly) {
     ASSERT_EQ(bucket4->getEra(), 2);
     ASSERT_EQ(_eraManager.getCountForEra(2), 1);
     clear(ns2);
+    checkAndRemoveClearedBucket(bucket4, bucketKey2, withLock);
     ASSERT_EQ(_eraManager.getCountForEra(2), 0);
     ASSERT_EQ(_eraManager.getCountForEra(1), 1);
 }
@@ -220,56 +233,65 @@ TEST_F(BucketCatalogEraTest, ClearRegistryGarbageCollection) {
     ASSERT_EQ(bucket2->getEra(), 0);
     ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 0);
     clear(ns1);
+    checkAndRemoveClearedBucket(bucket1, bucketKey1, withLock);
     // Era 0 still has non-zero count after this clear because bucket2 is still in era 0.
     ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
     clear(ns2);
+    checkAndRemoveClearedBucket(bucket2, bucketKey2, withLock);
     // Bucket2 gets deleted, which makes era 0's count decrease to 0, then clear registry gets
-    // cleaned (which removes the clear with era 1), and then inserts the new clear with era 2 into
-    // the registry.
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
+    // cleaned.
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 0);
 
     auto bucket3 = createBucket(info1);
     auto bucket4 = createBucket(info2);
     ASSERT_EQ(bucket3->getEra(), 2);
     ASSERT_EQ(bucket4->getEra(), 2);
     clear(ns1);
+    checkAndRemoveClearedBucket(bucket3, bucketKey1, withLock);
     // Era 2 still has bucket4 in it, so its count remains non-zero.
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 2);
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
     auto bucket5 = createBucket(info1);
     auto bucket6 = createBucket(info2);
     ASSERT_EQ(bucket5->getEra(), 3);
     ASSERT_EQ(bucket6->getEra(), 3);
     clear(ns1);
+    checkAndRemoveClearedBucket(bucket5, bucketKey1, withLock);
     // Eras 2 and 3 still have bucket4 and bucket6 in them respectively, so their counts remain
     // non-zero.
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 3);
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 2);
     clear(ns2);
+    checkAndRemoveClearedBucket(bucket4, bucketKey2, withLock);
+    checkAndRemoveClearedBucket(bucket6, bucketKey2, withLock);
     // Eras 2 and 3 have their counts become 0 because bucket4 and bucket6 are cleared. The clear
-    // ops with eras 2-4 are removed, and then clear op with era 5 is inserted.
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
+    // registry is emptied.
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 0);
 
     auto bucket7 = createBucket(info1);
     auto bucket8 = createBucket(info3);
     ASSERT_EQ(bucket7->getEra(), 5);
     ASSERT_EQ(bucket8->getEra(), 5);
     clear(ns3);
+    checkAndRemoveClearedBucket(bucket8, bucketKey3, withLock);
     // Era 5 still has bucket7 in it so its count remains non-zero.
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 2);
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
     auto bucket9 = createBucket(info2);
     ASSERT_EQ(bucket9->getEra(), 6);
     clear(ns2);
-    // Era 6's count becomes 0. The clear op with era 5 is removed from the clear registry.
+    checkAndRemoveClearedBucket(bucket9, bucketKey2, withLock);
+    // Era 6's count becomes 0. Since era 5 is the smallest era with non-zero count, no clear ops
+    // are removed.
     ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 2);
     auto bucket10 = createBucket(info3);
     ASSERT_EQ(bucket10->getEra(), 7);
     clear(ns3);
+    checkAndRemoveClearedBucket(bucket10, bucketKey3, withLock);
     // Era 7's count becomes 0. Since era 5 is the smallest era with non-zero count, no clear ops
     // are removed.
     ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 3);
     clear(ns1);
-    // Era 5's count becomes 0. No eras with non-zero counts remain, so all clear ops are removed
-    // (namely, clear ops with eras 6 and 7).
-    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 1);
+    checkAndRemoveClearedBucket(bucket7, bucketKey1, withLock);
+    // Era 5's count becomes 0. No eras with non-zero counts remain, so all clear ops are removed.
+    ASSERT_EQUALS(_eraManager.getClearOperationsCount(), 0);
 }