SERVER-68571 Update reaper when instantiating collection/index on expired ident

4 files changed, 162 insertions, 25 deletions

diff --git a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.cpp b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.cpp
index cc4b0acfcb0..03dd8d2ede3 100644
--- a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.cpp
+++ b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.cpp
@@ -54,14 +54,17 @@ void KVDropPendingIdentReaper::addDropPendingIdent(const Timestamp& dropTimestam
     const auto equalRange = _dropPendingIdents.equal_range(dropTimestamp);
     const auto& lowerBound = equalRange.first;
     const auto& upperBound = equalRange.second;
-    auto matcher = [ident](const auto& pair) { return pair.second.identName == ident->getIdent(); };
+    auto matcher = [ident](const auto& pair) {
+        return pair.second->identName == ident->getIdent();
+    };
     if (std::find_if(lowerBound, upperBound, matcher) == upperBound) {
-        IdentInfo info;
-        info.identName = ident->getIdent();
-        info.isDropped = false;
-        info.dropToken = ident;
-        info.onDrop = std::move(onDrop);
+        auto info = std::make_shared<IdentInfo>();
+        info->identName = ident->getIdent();
+        info->identState = IdentInfo::State::kNotDropped;
+        info->dropToken = ident;
+        info->onDrop = std::move(onDrop);
         _dropPendingIdents.insert(std::make_pair(dropTimestamp, info));
+        _identToTimestamp.insert(std::make_pair(ident->getIdent(), dropTimestamp));
     } else {
         LOGV2_FATAL_NOTRACE(51023,
                             "Failed to add drop-pending ident, duplicate timestamp and ident pair",
@@ -70,6 +73,44 @@ void KVDropPendingIdentReaper::addDropPendingIdent(const Timestamp& dropTimestam
     }
 }
 
+std::shared_ptr<Ident> KVDropPendingIdentReaper::markIdentInUse(const std::string& ident) {
+    stdx::lock_guard<Latch> lock(_mutex);
+    auto identToTimestampIt = _identToTimestamp.find(ident);
+    if (identToTimestampIt == _identToTimestamp.end()) {
+        // Ident is not known to the reaper.
+        return nullptr;
+    }
+
+    auto beginEndPair = _dropPendingIdents.equal_range(identToTimestampIt->second);
+    for (auto dropPendingIdentsIt = beginEndPair.first; dropPendingIdentsIt != beginEndPair.second;
+         dropPendingIdentsIt++) {
+        auto info = dropPendingIdentsIt->second;
+        if (info->identName != ident) {
+            continue;
+        }
+
+        if (info->identState == IdentInfo::State::kBeingDropped ||
+            info->identState == IdentInfo::State::kDropped) {
+            // The ident is being dropped or was already dropped. Cannot mark the ident as in use.
+            return nullptr;
+        }
+
+        if (auto existingIdent = info->dropToken.lock()) {
+            // This function can be called concurrently and we need to share the same ident at any
+            // given time to prevent the reaper from removing idents prematurely.
+            return existingIdent;
+        }
+
+        std::shared_ptr<Ident> newIdent = std::make_shared<Ident>(info->identName);
+        info->dropToken = newIdent;
+        return newIdent;
+    }
+
+    // The ident was found in _identToTimestamp earlier, so it must exist in _dropPendingIdents.
+    MONGO_UNREACHABLE;
+}
+
+
 boost::optional<Timestamp> KVDropPendingIdentReaper::getEarliestDropTimestamp() const {
     stdx::lock_guard<Latch> lock(_mutex);
     auto it = _dropPendingIdents.cbegin();
@@ -84,7 +125,7 @@ std::set<std::string> KVDropPendingIdentReaper::getAllIdentNames() const {
     std::set<std::string> identNames;
     for (const auto& entry : _dropPendingIdents) {
         const auto& identInfo = entry.second;
-        const auto& identName = identInfo.identName;
+        const auto& identName = identInfo->identName;
         identNames.insert(identName);
     }
     return identNames;
@@ -94,13 +135,15 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
     DropPendingIdents toDrop;
     {
         stdx::lock_guard<Latch> lock(_mutex);
-        for (auto it = _dropPendingIdents.cbegin();
-             it != _dropPendingIdents.cend() && (it->first < ts || it->first == Timestamp::min());
+        for (auto it = _dropPendingIdents.begin();
+             it != _dropPendingIdents.end() && (it->first < ts || it->first == Timestamp::min());
              ++it) {
             // This collection/index satisfies the 'ts' requirement to be safe to drop, but we must
             // also check that there are no active operations remaining that still retain a
             // reference by which to access the collection/index data.
-            if (it->second.dropToken.expired()) {
+            const auto& info = it->second;
+            if (info->identState == IdentInfo::State::kNotDropped && info->dropToken.expired()) {
+                info->identState = IdentInfo::State::kBeingDropped;
                 toDrop.insert(*it);
             }
         }
@@ -118,14 +161,13 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
 
             const auto& dropTimestamp = timestampAndIdentInfo.first;
             auto& identInfo = timestampAndIdentInfo.second;
-            const auto& identName = identInfo.identName;
+            const auto& identName = identInfo->identName;
             LOGV2(22237,
                   "Completing drop for ident",
                   "ident"_attr = identName,
                   "dropTimestamp"_attr = dropTimestamp);
             WriteUnitOfWork wuow(opCtx);
-            auto status =
-                _engine->dropIdent(opCtx->recoveryUnit(), identName, std::move(identInfo.onDrop));
+            auto status = _engine->dropIdent(opCtx->recoveryUnit(), identName, identInfo->onDrop);
             if (!status.isOK()) {
                 if (status == ErrorCodes::ObjectIsBusy) {
                     LOGV2(6936300,
@@ -133,6 +175,9 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
                           "ident"_attr = identName,
                           "dropTimestamp"_attr = dropTimestamp,
                           "error"_attr = status);
+
+                    stdx::lock_guard<Latch> lock(_mutex);
+                    identInfo->identState = IdentInfo::State::kNotDropped;
                     return;
                 }
                 LOGV2_FATAL_NOTRACE(51022,
@@ -142,9 +187,12 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
                                     "error"_attr = status);
             }
 
-            // Ident drops are non-transactional and cannot be rolled back. So this does not need to
-            // be in an onCommit handler.
-            identInfo.isDropped = true;
+            {
+                // Ident drops are non-transactional and cannot be rolled back. So this does not
+                // need to be in an onCommit handler.
+                stdx::lock_guard<Latch> lock(_mutex);
+                identInfo->identState = IdentInfo::State::kDropped;
+            }
 
             wuow.commit();
             LOGV2(6776600,
@@ -160,7 +208,10 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
 
         stdx::lock_guard<Latch> lock(_mutex);
         for (const auto& timestampAndIdentInfo : toDrop) {
-            if (!timestampAndIdentInfo.second.isDropped) {
+            // The ident was either dropped or put back in a not dropped state.
+            invariant(timestampAndIdentInfo.second->identState != IdentInfo::State::kBeingDropped);
+
+            if (timestampAndIdentInfo.second->identState == IdentInfo::State::kNotDropped) {
                 // This ident was not dropped. Skip removing it from the drop pending list.
                 continue;
             }
@@ -171,7 +222,8 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
             // ident as well as the timestamp.
             auto beginEndPair = _dropPendingIdents.equal_range(timestampAndIdentInfo.first);
             for (auto it = beginEndPair.first; it != beginEndPair.second;) {
-                if (it->second.identName == timestampAndIdentInfo.second.identName) {
+                if (it->second == timestampAndIdentInfo.second) {
+                    invariant(_identToTimestamp.erase(it->second->identName) == 1);
                     it = _dropPendingIdents.erase(it);
                     break;
                 } else {
@@ -185,6 +237,7 @@ void KVDropPendingIdentReaper::dropIdentsOlderThan(OperationContext* opCtx, cons
 void KVDropPendingIdentReaper::clearDropPendingState() {
     stdx::lock_guard<Latch> lock(_mutex);
     _dropPendingIdents.clear();
+    _identToTimestamp.clear();
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.h b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.h
index 08095e05201..e0497f07cb0 100644
--- a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.h
+++ b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper.h
@@ -84,6 +84,15 @@ public:
                              StorageEngine::DropIdentCallback&& onDrop = nullptr);
 
     /**
+     * Marks the ident as in use and prevents the reaper from dropping the ident.
+     *
+     * Returns nullptr if the ident is not found, or if the ident state is `kBeingDropped` or
+     * `kDropped`. Returns a shared_ptr to the `dropToken` if it isn't expired, otherwise a new
+     * shared_ptr is generated, stored in `dropToken`, and returned.
+     */
+    std::shared_ptr<Ident> markIdentInUse(const std::string& ident);
+
+    /**
      * Returns earliest drop timestamp in '_dropPendingIdents'.
      * Returns boost::none if '_dropPendingIdents' is empty.
      */
@@ -115,12 +124,13 @@ private:
         // Identifier for the storage to drop the associated collection or index data.
         std::string identName;
 
-        // Whether the storage engine dropped the ident.
-        bool isDropped;
+        // Ident drop state.
+        enum class State { kNotDropped, kBeingDropped, kDropped };
+        State identState;
 
         // The collection or index data can be safely dropped when no references to this token
         // remain.
-        std::weak_ptr<void> dropToken;
+        std::weak_ptr<Ident> dropToken;
 
         // Callback to run once the ident has been dropped.
         StorageEngine::DropIdentCallback onDrop;
@@ -130,7 +140,7 @@ private:
     // namespaces by drop optime. Additionally, it is possible for certain user operations (such
     // as renameCollection across databases) to generate more than one drop-pending namespace for
     // the same drop optime.
-    using DropPendingIdents = std::multimap<Timestamp, IdentInfo>;
+    using DropPendingIdents = std::multimap<Timestamp, std::shared_ptr<IdentInfo>>;
 
     // Used to access the KV engine for the purposes of dropping the ident.
     KVEngine* const _engine;
@@ -140,6 +150,9 @@ private:
 
     // Drop-pending idents. Ordered by drop timestamp.
     DropPendingIdents _dropPendingIdents;
+
+    // Ident to drop timestamp map. Used for efficient lookups into _dropPendingIdents.
+    StringMap<Timestamp> _identToTimestamp;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
index 47797efdf48..215dbee3b46 100644
--- a/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
+++ b/src/mongo/db/storage/kv/kv_drop_pending_ident_reaper_test.cpp
@@ -48,7 +48,7 @@ class KVEngineMock : public KVEngine {
 public:
     Status dropIdent(RecoveryUnit* ru,
                      StringData ident,
-                     StorageEngine::DropIdentCallback&& onDrop) override;
+                     const StorageEngine::DropIdentCallback& onDrop) override;
 
     void dropIdentForImport(OperationContext* opCtx, StringData ident) override {}
 
@@ -160,7 +160,7 @@ public:
 
 Status KVEngineMock::dropIdent(RecoveryUnit* ru,
                                StringData ident,
-                               StorageEngine::DropIdentCallback&& onDrop) {
+                               const StorageEngine::DropIdentCallback& onDrop) {
     auto status = dropIdentFn(ru, ident);
     if (status.isOK()) {
         droppedIdents.push_back(ident.toString());
@@ -390,6 +390,77 @@ TEST_F(KVDropPendingIdentReaperTest, DropIdentsOlderThanSkipsIdentsStillReferenc
     ASSERT_FALSE(reaper.getEarliestDropTimestamp());
 }
 
+TEST_F(KVDropPendingIdentReaperTest, MarkUnknownIdentInUse) {
+    const std::string identName = "ident";
+    auto engine = getEngine();
+
+    KVDropPendingIdentReaper reaper(engine);
+
+    std::shared_ptr<Ident> newIdent = reaper.markIdentInUse(identName);
+    ASSERT(!newIdent);
+}
+
+TEST_F(KVDropPendingIdentReaperTest, MarkUnexpiredIdentInUse) {
+    const std::string identName = "ident";
+    const Timestamp dropTimestamp = Timestamp(50, 50);
+    auto engine = getEngine();
+
+    KVDropPendingIdentReaper reaper(engine);
+
+    // The reaper will not have an expired reference to the ident.
+    std::shared_ptr<Ident> ident = std::make_shared<Ident>(identName);
+    reaper.addDropPendingIdent(dropTimestamp, ident);
+
+    ASSERT_EQUALS(dropTimestamp, *reaper.getEarliestDropTimestamp());
+
+    // Marking an unexpired ident as in-use will return a shared_ptr to that ident.
+    std::shared_ptr<Ident> newIdent = reaper.markIdentInUse(identName);
+    ASSERT_EQ(ident, newIdent);
+    ASSERT_EQ(2, ident.use_count());
+
+    auto opCtx = makeOpCtx();
+    reaper.dropIdentsOlderThan(opCtx.get(), Timestamp::max());
+    ASSERT_EQUALS(0U, engine->droppedIdents.size());
+
+    // Remove the references to the ident so that the reaper can drop it the next time.
+    ident.reset();
+    newIdent.reset();
+
+    reaper.dropIdentsOlderThan(opCtx.get(), Timestamp::max());
+    ASSERT_EQUALS(1U, engine->droppedIdents.size());
+    ASSERT_EQUALS(identName, engine->droppedIdents.front());
+}
+
+TEST_F(KVDropPendingIdentReaperTest, MarkExpiredIdentInUse) {
+    const std::string identName = "ident";
+    const Timestamp dropTimestamp = Timestamp(50, 50);
+    auto engine = getEngine();
+
+    KVDropPendingIdentReaper reaper(engine);
+    {
+        // The reaper will have an expired weak_ptr to the ident.
+        std::shared_ptr<Ident> ident = std::make_shared<Ident>(identName);
+        reaper.addDropPendingIdent(dropTimestamp, ident);
+    }
+
+    ASSERT_EQUALS(dropTimestamp, *reaper.getEarliestDropTimestamp());
+
+    // Mark the ident as in use to prevent the reaper from dropping it.
+    std::shared_ptr<Ident> ident = reaper.markIdentInUse(identName);
+    ASSERT_EQ(1, ident.use_count());
+
+    auto opCtx = makeOpCtx();
+    reaper.dropIdentsOlderThan(opCtx.get(), Timestamp::max());
+    ASSERT_EQUALS(0U, engine->droppedIdents.size());
+
+    // Remove the reference to the ident so that the reaper can drop it the next time.
+    ident.reset();
+
+    reaper.dropIdentsOlderThan(opCtx.get(), Timestamp::max());
+    ASSERT_EQUALS(1U, engine->droppedIdents.size());
+    ASSERT_EQUALS(identName, engine->droppedIdents.front());
+}
+
 DEATH_TEST_F(KVDropPendingIdentReaperTest,
              DropIdentsOlderThanTerminatesIfKVEngineFailsToDropIdent,
              "Failed to remove drop-pending ident") {
diff --git a/src/mongo/db/storage/kv/kv_engine.h b/src/mongo/db/storage/kv/kv_engine.h
index d7a1a66323a..557a568d3d3 100644
--- a/src/mongo/db/storage/kv/kv_engine.h
+++ b/src/mongo/db/storage/kv/kv_engine.h
@@ -164,7 +164,7 @@ public:
      */
     virtual Status dropIdent(RecoveryUnit* ru,
                              StringData ident,
-                             StorageEngine::DropIdentCallback&& onDrop = nullptr) = 0;
+                             const StorageEngine::DropIdentCallback& onDrop = nullptr) = 0;
 
     /**
      * Removes any knowledge of the ident from the storage engines metadata without removing the