SERVER-42834: Register an onRollback handler when adding a drop pending namespace.

5 files changed, 51 insertions, 29 deletions

diff --git a/src/mongo/db/repl/drop_pending_collection_reaper.cpp b/src/mongo/db/repl/drop_pending_collection_reaper.cpp
index 31993c6acc2..cb4c85c5cd4 100644
--- a/src/mongo/db/repl/drop_pending_collection_reaper.cpp
+++ b/src/mongo/db/repl/drop_pending_collection_reaper.cpp
@@ -76,7 +76,9 @@ DropPendingCollectionReaper::DropPendingCollectionReaper(StorageInterface* stora
     : _storageInterface(storageInterface) {}
 
 void DropPendingCollectionReaper::addDropPendingNamespace(
-    const OpTime& dropOpTime, const NamespaceString& dropPendingNamespace) {
+    OperationContext* opCtx,
+    const OpTime& dropOpTime,
+    const NamespaceString& dropPendingNamespace) {
     invariant(dropPendingNamespace.isDropPendingNamespace());
     stdx::lock_guard<stdx::mutex> lock(_mutex);
     const auto equalRange = _dropPendingNamespaces.equal_range(dropOpTime);
@@ -85,13 +87,30 @@ void DropPendingCollectionReaper::addDropPendingNamespace(
     auto matcher = [&dropPendingNamespace](const auto& pair) {
         return pair.second == dropPendingNamespace;
     };
-    if (std::find_if(lowerBound, upperBound, matcher) == upperBound) {
-        _dropPendingNamespaces.insert(std::make_pair(dropOpTime, dropPendingNamespace));
-    } else {
+
+    if (std::find_if(lowerBound, upperBound, matcher) != upperBound) {
         severe() << "Failed to add drop-pending collection " << dropPendingNamespace
                  << " with drop optime " << dropOpTime << ": duplicate optime and namespace pair.";
         fassertFailedNoTrace(40448);
     }
+
+    _dropPendingNamespaces.insert(std::make_pair(dropOpTime, dropPendingNamespace));
+    if (opCtx->lockState()->inAWriteUnitOfWork()) {
+        opCtx->recoveryUnit()->onRollback([this, dropPendingNamespace, dropOpTime]() {
+            stdx::lock_guard<stdx::mutex> lock(_mutex);
+
+            const auto equalRange = _dropPendingNamespaces.equal_range(dropOpTime);
+            const auto& lowerBound = equalRange.first;
+            const auto& upperBound = equalRange.second;
+            auto matcher = [&dropPendingNamespace](const auto& pair) {
+                return pair.second == dropPendingNamespace;
+            };
+
+            auto it = std::find_if(lowerBound, upperBound, matcher);
+            invariant(it != upperBound);
+            _dropPendingNamespaces.erase(it);
+        });
+    }
 }
 
 boost::optional<OpTime> DropPendingCollectionReaper::getEarliestDropOpTime() {
diff --git a/src/mongo/db/repl/drop_pending_collection_reaper.h b/src/mongo/db/repl/drop_pending_collection_reaper.h
index be8dd9a77d8..48795159066 100644
--- a/src/mongo/db/repl/drop_pending_collection_reaper.h
+++ b/src/mongo/db/repl/drop_pending_collection_reaper.h
@@ -82,7 +82,8 @@ public:
     /**
      * Adds a new drop-pending namespace, with its drop optime, to be managed by this class.
      */
-    void addDropPendingNamespace(const OpTime& dropOpTime,
+    void addDropPendingNamespace(OperationContext* opCtx,
+                                 const OpTime& dropOpTime,
                                  const NamespaceString& dropPendingNamespace);
 
     /**
diff --git a/src/mongo/db/repl/drop_pending_collection_reaper_test.cpp b/src/mongo/db/repl/drop_pending_collection_reaper_test.cpp
index 43fb891728f..0ac9f2071e6 100644
--- a/src/mongo/db/repl/drop_pending_collection_reaper_test.cpp
+++ b/src/mongo/db/repl/drop_pending_collection_reaper_test.cpp
@@ -118,7 +118,7 @@ TEST_F(DropPendingCollectionReaperTest, AddDropPendingNamespaceAcceptsNullDropOp
     OpTime nullDropOpTime;
     auto dpns = NamespaceString("test.foo").makeDropPendingNamespace(nullDropOpTime);
     DropPendingCollectionReaper reaper(_storageInterface.get());
-    reaper.addDropPendingNamespace(nullDropOpTime, dpns);
+    reaper.addDropPendingNamespace(makeOpCtx().get(), nullDropOpTime, dpns);
     ASSERT_EQUALS(nullDropOpTime, *reaper.getEarliestDropOpTime());
 }
 
@@ -134,12 +134,12 @@ TEST_F(DropPendingCollectionReaperTest,
 
     OpTime opTime({Seconds(100), 0}, 1LL);
     auto dpns = NamespaceString("test.foo").makeDropPendingNamespace(opTime);
-    reaper.addDropPendingNamespace(opTime, dpns);
-    reaper.addDropPendingNamespace(opTime,
-                                   NamespaceString("test.bar").makeDropPendingNamespace(opTime));
+    auto opCtx = makeOpCtx();
+    reaper.addDropPendingNamespace(opCtx.get(), opTime, dpns);
+    reaper.addDropPendingNamespace(
+        opCtx.get(), opTime, NamespaceString("test.bar").makeDropPendingNamespace(opTime));
 
     // Drop all collections managed by reaper and confirm number of drops.
-    auto opCtx = makeOpCtx();
     reaper.dropCollectionsOlderThan(opCtx.get(), opTime);
     ASSERT_EQUALS(2U, numCollectionsDropped);
 }
@@ -150,8 +150,9 @@ DEATH_TEST_F(DropPendingCollectionReaperTest,
     OpTime opTime({Seconds(100), 0}, 1LL);
     auto dpns = NamespaceString("test.foo").makeDropPendingNamespace(opTime);
     DropPendingCollectionReaper reaper(_storageInterface.get());
-    reaper.addDropPendingNamespace(opTime, dpns);
-    reaper.addDropPendingNamespace(opTime, dpns);
+    auto opCtx = makeOpCtx();
+    reaper.addDropPendingNamespace(opCtx.get(), opTime, dpns);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime, dpns);
 }
 
 TEST_F(DropPendingCollectionReaperTest,
@@ -176,11 +177,11 @@ TEST_F(DropPendingCollectionReaperTest,
     // getEarliestDropOpTime() returns earliest optime.
     DropPendingCollectionReaper reaper(_storageInterface.get());
     ASSERT_FALSE(reaper.getEarliestDropOpTime());
-    reaper.addDropPendingNamespace(opTime[1], dpns[1]);
-    reaper.addDropPendingNamespace(opTime[0], dpns[0]);
-    reaper.addDropPendingNamespace(opTime[2], dpns[2]);
-    reaper.addDropPendingNamespace(opTime[3], dpns[3]);
-    reaper.addDropPendingNamespace(opTime[4], dpns[4]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[1], dpns[1]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[0], dpns[0]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[2], dpns[2]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[3], dpns[3]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[4], dpns[4]);
     ASSERT_EQUALS(opTime[0], *reaper.getEarliestDropOpTime());
 
     // Committed optime before first drop optime has no effect.
@@ -215,9 +216,9 @@ TEST_F(DropPendingCollectionReaperTest, DropCollectionsOlderThanHasNoEffectIfCol
     auto dpns = ns.makeDropPendingNamespace(optime);
 
     DropPendingCollectionReaper reaper(_storageInterface.get());
-    reaper.addDropPendingNamespace(optime, dpns);
 
     auto opCtx = makeOpCtx();
+    reaper.addDropPendingNamespace(opCtx.get(), optime, dpns);
     reaper.dropCollectionsOlderThan(opCtx.get(), optime);
 }
 
@@ -230,9 +231,9 @@ TEST_F(DropPendingCollectionReaperTest, DropCollectionsOlderThanLogsDropCollecti
     StorageInterfaceMock storageInterfaceMock;
 
     DropPendingCollectionReaper reaper(&storageInterfaceMock);
-    reaper.addDropPendingNamespace(optime, dpns);
-
     auto opCtx = makeOpCtx();
+
+    reaper.addDropPendingNamespace(opCtx.get(), optime, dpns);
     startCapturingLogMessages();
     reaper.dropCollectionsOlderThan(opCtx.get(), optime);
     stopCapturingLogMessages();
@@ -259,9 +260,9 @@ TEST_F(DropPendingCollectionReaperTest,
     };
 
     DropPendingCollectionReaper reaper(&storageInterfaceMock);
-    reaper.addDropPendingNamespace(optime, dpns);
 
     auto opCtx = makeOpCtx();
+    reaper.addDropPendingNamespace(opCtx.get(), optime, dpns);
     reaper.dropCollectionsOlderThan(opCtx.get(), optime);
 
     ASSERT_EQUALS(dpns, droppedNss);
@@ -286,9 +287,9 @@ TEST_F(DropPendingCollectionReaperTest, RollBackDropPendingCollection) {
     }
 
     DropPendingCollectionReaper reaper(_storageInterface.get());
-    reaper.addDropPendingNamespace(opTime[0], dpns[0]);
-    reaper.addDropPendingNamespace(opTime[1], dpns[1]);
-    reaper.addDropPendingNamespace(opTime[2], dpns[2]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[0], dpns[0]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[1], dpns[1]);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[2], dpns[2]);
 
     // Rolling back at an optime not in the list returns false.
     ASSERT_FALSE(
@@ -318,7 +319,7 @@ TEST_F(DropPendingCollectionReaperTest, RollBackDropPendingCollection) {
     NamespaceString ns4 = NamespaceString("test", "coll4");
     NamespaceString dpns4 = ns4.makeDropPendingNamespace(opTime[1]);
     ASSERT_OK(_storageInterface->createCollection(opCtx.get(), dpns4, generateOptionsWithUuid()));
-    reaper.addDropPendingNamespace(opTime[1], dpns4);
+    reaper.addDropPendingNamespace(opCtx.get(), opTime[1], dpns4);
     ASSERT_TRUE(reaper.rollBackDropPendingCollection(opCtx.get(), opTime[1], ns[1]));
     ASSERT_EQUALS(opTime[1], *reaper.getEarliestDropOpTime());
     ASSERT_TRUE(collectionExists(opCtx.get(), dpns[0]));
diff --git a/src/mongo/db/repl/rollback_impl_test.cpp b/src/mongo/db/repl/rollback_impl_test.cpp
index 05ce9e82cf9..e3cb9c970c9 100644
--- a/src/mongo/db/repl/rollback_impl_test.cpp
+++ b/src/mongo/db/repl/rollback_impl_test.cpp
@@ -1138,7 +1138,8 @@ TEST_F(RollbackImplTest, RollbackProperlySavesFilesWhenInsertsAndDropOfCollectio
     const auto nss = NamespaceString("db.people").makeDropPendingNamespace(dropOpTime);
     const auto uuid = UUID::gen();
     const auto coll = _initializeCollection(_opCtx.get(), uuid, nss);
-    DropPendingCollectionReaper::get(_opCtx.get())->addDropPendingNamespace(dropOpTime, nss);
+    DropPendingCollectionReaper::get(_opCtx.get())
+        ->addDropPendingNamespace(_opCtx.get(), dropOpTime, nss);
 
     // Insert documents into the collection. We'll write them out even though the collection is
     // later dropped.
diff --git a/src/mongo/db/repl/rs_rollback_test.cpp b/src/mongo/db/repl/rs_rollback_test.cpp
index 6d1bfbda79a..6408443af27 100644
--- a/src/mongo/db/repl/rs_rollback_test.cpp
+++ b/src/mongo/db/repl/rs_rollback_test.cpp
@@ -964,7 +964,7 @@ TEST_F(RSRollbackTest, RollbackDropCollectionCommand) {
     CollectionOptions options;
     options.uuid = UUID::gen();
     auto coll = _createCollection(_opCtx.get(), dpns, options);
-    _dropPendingCollectionReaper->addDropPendingNamespace(dropTime, dpns);
+    _dropPendingCollectionReaper->addDropPendingNamespace(_opCtx.get(), dropTime, dpns);
 
     auto commonOperation = makeOpAndRecordId(1);
     auto dropCollectionOperation =
@@ -1126,7 +1126,7 @@ TEST_F(RSRollbackTest, RollbackRenameCollectionInDatabaseWithDropTargetTrueComma
     CollectionOptions droppedCollOptions;
     droppedCollOptions.uuid = UUID::gen();
     auto droppedColl = _createCollection(_opCtx.get(), dpns, droppedCollOptions);
-    _dropPendingCollectionReaper->addDropPendingNamespace(dropTime, dpns);
+    _dropPendingCollectionReaper->addDropPendingNamespace(_opCtx.get(), dropTime, dpns);
     auto droppedCollectionUUID = droppedColl->uuid();
 
     CollectionOptions renamedCollOptions;
@@ -1285,7 +1285,7 @@ TEST_F(RSRollbackTest, RollbackDropCollectionThenRenameCollectionToDroppedCollec
     droppedCollOptions.uuid = UUID::gen();
     auto droppedCollection = _createCollection(_opCtx.get(), dpns, droppedCollOptions);
     auto droppedCollectionUUID = droppedCollection->uuid();
-    _dropPendingCollectionReaper->addDropPendingNamespace(dropTime, dpns);
+    _dropPendingCollectionReaper->addDropPendingNamespace(_opCtx.get(), dropTime, dpns);
 
     auto commonOperation = makeOpAndRecordId(1);