SERVER-37726 Add the ability to abort an index builder by the index names it is building

4 files changed, 191 insertions, 42 deletions

diff --git a/src/mongo/db/catalog/index_builds_manager.cpp b/src/mongo/db/catalog/index_builds_manager.cpp
index 6e88a1985c0..62d1e2c8678 100644
--- a/src/mongo/db/catalog/index_builds_manager.cpp
+++ b/src/mongo/db/catalog/index_builds_manager.cpp
@@ -89,7 +89,7 @@ Status IndexBuildsManager::setUpIndexBuild(OperationContext* opCtx,
               str::stream() << "Unable to set up index build " << buildUUID << ": collection "
                             << nss.ns() << " is not locked in exclusive mode.");
 
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
     if (options.protocol == IndexBuildProtocol::kTwoPhase) {
         builder->setTwoPhaseBuildUUID(buildUUID);
     }
@@ -124,14 +124,14 @@ Status IndexBuildsManager::setUpIndexBuild(OperationContext* opCtx,
 Status IndexBuildsManager::startBuildingIndex(OperationContext* opCtx,
                                               Collection* collection,
                                               const UUID& buildUUID) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
 
     return builder->insertAllDocumentsInCollection(opCtx, collection);
 }
 
 StatusWith<std::pair<long long, long long>> IndexBuildsManager::startBuildingIndexForRecovery(
     OperationContext* opCtx, NamespaceString ns, const UUID& buildUUID, RepairData repair) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
 
     auto coll = CollectionCatalog::get(opCtx).lookupCollectionByNamespace(opCtx, ns);
     auto rs = coll ? coll->getRecordStore() : nullptr;
@@ -214,7 +214,7 @@ Status IndexBuildsManager::drainBackgroundWrites(
     const UUID& buildUUID,
     RecoveryUnit::ReadSource readSource,
     IndexBuildInterceptor::DrainYieldPolicy drainYieldPolicy) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
 
     return builder->drainBackgroundWrites(opCtx, readSource, drainYieldPolicy);
 }
@@ -222,13 +222,13 @@ Status IndexBuildsManager::drainBackgroundWrites(
 Status IndexBuildsManager::retrySkippedRecords(OperationContext* opCtx,
                                                const UUID& buildUUID,
                                                Collection* collection) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
     return builder->retrySkippedRecords(opCtx, collection);
 }
 
 Status IndexBuildsManager::checkIndexConstraintViolations(OperationContext* opCtx,
                                                           const UUID& buildUUID) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
 
     return builder->checkConstraints(opCtx);
 }
@@ -239,13 +239,13 @@ Status IndexBuildsManager::commitIndexBuild(OperationContext* opCtx,
                                             const UUID& buildUUID,
                                             MultiIndexBlock::OnCreateEachFn onCreateEachFn,
                                             MultiIndexBlock::OnCommitFn onCommitFn) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
 
     return writeConflictRetry(
         opCtx,
         "IndexBuildsManager::commitIndexBuild",
         nss.ns(),
-        [builder, opCtx, collection, nss, &onCreateEachFn, &onCommitFn] {
+        [this, builder, buildUUID, opCtx, collection, nss, &onCreateEachFn, &onCommitFn] {
             WriteUnitOfWork wunit(opCtx);
             auto status = builder->commit(opCtx, collection, onCreateEachFn, onCommitFn);
             if (!status.isOK()) {
@@ -258,6 +258,10 @@ Status IndexBuildsManager::commitIndexBuild(OperationContext* opCtx,
             // IndexBuilder. See SERVER-38986 and SERVER-34896.
             IndexTimestampHelper::setGhostCommitTimestampForCatalogWrite(opCtx, nss);
             wunit.commit();
+
+            // Required call to clean up even though commit cleaned everything up.
+            builder->cleanUpAfterBuild(opCtx, collection, MultiIndexBlock::kNoopOnCleanUpFn);
+            _unregisterIndexBuild(buildUUID);
             return Status::OK();
         });
 }
@@ -278,12 +282,11 @@ bool IndexBuildsManager::abortIndexBuild(const UUID& buildUUID, const std::strin
 }
 
 bool IndexBuildsManager::abortIndexBuildWithoutCleanup(OperationContext* opCtx,
+                                                       Collection* collection,
                                                        const UUID& buildUUID,
                                                        const std::string& reason) {
-    stdx::unique_lock<Latch> lk(_mutex);
-
-    auto builderIt = _builders.find(buildUUID);
-    if (builderIt == _builders.end()) {
+    auto builder = _getBuilder(buildUUID);
+    if (!builder.isOK()) {
         return false;
     }
 
@@ -291,10 +294,10 @@ bool IndexBuildsManager::abortIndexBuildWithoutCleanup(OperationContext* opCtx,
           "Index build aborted without cleanup: {buildUUID}: {reason}",
           "buildUUID"_attr = buildUUID,
           "reason"_attr = reason);
-    std::shared_ptr<MultiIndexBlock> builder = builderIt->second;
 
-    lk.unlock();
-    builder->abortWithoutCleanup(opCtx);
+    builder.getValue()->abortWithoutCleanup(opCtx);
+    builder.getValue()->cleanUpAfterBuild(opCtx, collection, MultiIndexBlock::kNoopOnCleanUpFn);
+    _unregisterIndexBuild(buildUUID);
 
     return true;
 }
@@ -303,14 +306,17 @@ void IndexBuildsManager::tearDownIndexBuild(OperationContext* opCtx,
                                             Collection* collection,
                                             const UUID& buildUUID,
                                             OnCleanUpFn onCleanUpFn) {
-    // TODO verify that the index builder is in a finished state before allowing its destruction.
     auto builder = _getBuilder(buildUUID);
-    builder->cleanUpAfterBuild(opCtx, collection, onCleanUpFn);
+    if (!builder.isOK()) {
+        return;
+    }
+
+    builder.getValue()->cleanUpAfterBuild(opCtx, collection, onCleanUpFn);
     _unregisterIndexBuild(buildUUID);
 }
 
 bool IndexBuildsManager::isBackgroundBuilding(const UUID& buildUUID) {
-    auto builder = _getBuilder(buildUUID);
+    auto builder = invariant(_getBuilder(buildUUID));
     return builder->isBackgroundBuilding();
 }
 
@@ -329,14 +335,19 @@ void IndexBuildsManager::_unregisterIndexBuild(const UUID& buildUUID) {
     stdx::unique_lock<Latch> lk(_mutex);
 
     auto builderIt = _builders.find(buildUUID);
-    invariant(builderIt != _builders.end());
+    if (builderIt == _builders.end()) {
+        return;
+    }
     _builders.erase(builderIt);
 }
 
-std::shared_ptr<MultiIndexBlock> IndexBuildsManager::_getBuilder(const UUID& buildUUID) {
+StatusWith<std::shared_ptr<MultiIndexBlock>> IndexBuildsManager::_getBuilder(
+    const UUID& buildUUID) {
     stdx::unique_lock<Latch> lk(_mutex);
     auto builderIt = _builders.find(buildUUID);
-    invariant(builderIt != _builders.end());
+    if (builderIt == _builders.end()) {
+        return {ErrorCodes::NoSuchKey, str::stream() << "No index build with UUID: " << buildUUID};
+    }
     return builderIt->second;
 }
 
diff --git a/src/mongo/db/catalog/index_builds_manager.h b/src/mongo/db/catalog/index_builds_manager.h
index ad89d608d5b..87cada4d79f 100644
--- a/src/mongo/db/catalog/index_builds_manager.h
+++ b/src/mongo/db/catalog/index_builds_manager.h
@@ -153,6 +153,7 @@ public:
      * been cleared away, or not having yet started..
      */
     bool abortIndexBuildWithoutCleanup(OperationContext* opCtx,
+                                       Collection* collection,
                                        const UUID& buildUUID,
                                        const std::string& reason);
 
@@ -188,9 +189,9 @@ private:
     void _unregisterIndexBuild(const UUID& buildUUID);
 
     /**
-     * Returns a shared pointer to the builder. Invariants if the builder does not exist.
+     * Returns a shared pointer to the builder. Returns a bad status if the builder does not exist.
      */
-    std::shared_ptr<MultiIndexBlock> _getBuilder(const UUID& buildUUID);
+    StatusWith<std::shared_ptr<MultiIndexBlock>> _getBuilder(const UUID& buildUUID);
 
     // Protects the map data structures below.
     mutable Mutex _mutex = MONGO_MAKE_LATCH("IndexBuildsManager::_mutex");
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index 458921595da..4d71a91c8a5 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -508,25 +508,44 @@ void IndexBuildsCoordinator::waitForAllIndexBuildsToStopForShutdown() {
     }
 }
 
-void IndexBuildsCoordinator::abortCollectionIndexBuilds(const UUID& collectionUUID,
-                                                        const std::string& reason) {
-    stdx::unique_lock<Latch> lk(_mutex);
-
-    // Ensure the caller correctly stopped any new index builds on the collection.
-    auto it = _disallowedCollections.find(collectionUUID);
-    invariant(it != _disallowedCollections.end());
-
+std::vector<UUID> IndexBuildsCoordinator::_abortCollectionIndexBuilds(stdx::unique_lock<Latch>& lk,
+                                                                      const UUID& collectionUUID,
+                                                                      const std::string& reason,
+                                                                      bool shouldWait) {
     auto collIndexBuildsIt = _collectionIndexBuilds.find(collectionUUID);
     if (collIndexBuildsIt == _collectionIndexBuilds.end()) {
-        return;
+        return {};
     }
 
+    log() << "About to abort all index builders on collection with UUID: " << collectionUUID;
+
+    std::vector<UUID> buildUUIDs = collIndexBuildsIt->second->getIndexBuildUUIDs(lk);
     collIndexBuildsIt->second->runOperationOnAllBuilds(
         lk, &_indexBuildsManager, abortIndexBuild, reason);
+
+    if (!shouldWait) {
+        return buildUUIDs;
+    }
+
     // Take a shared ptr, rather than accessing the Tracker through the map's iterator, so that the
     // object does not destruct while we are waiting, causing a use-after-free memory error.
     auto collIndexBuildsSharedPtr = collIndexBuildsIt->second;
     collIndexBuildsSharedPtr->waitUntilNoIndexBuildsRemain(lk);
+    return buildUUIDs;
+}
+
+void IndexBuildsCoordinator::abortCollectionIndexBuilds(const UUID& collectionUUID,
+                                                        const std::string& reason) {
+    stdx::unique_lock<Latch> lk(_mutex);
+    const bool shouldWait = true;
+    _abortCollectionIndexBuilds(lk, collectionUUID, reason, shouldWait);
+}
+
+std::vector<UUID> IndexBuildsCoordinator::abortCollectionIndexBuildsNoWait(
+    const UUID& collectionUUID, const std::string& reason) {
+    stdx::unique_lock<Latch> lk(_mutex);
+    const bool shouldWait = false;
+    return _abortCollectionIndexBuilds(lk, collectionUUID, reason, shouldWait);
 }
 
 void IndexBuildsCoordinator::abortDatabaseIndexBuilds(StringData db, const std::string& reason) {
@@ -682,7 +701,8 @@ void IndexBuildsCoordinator::abortIndexBuildByBuildUUID(OperationContext* opCtx,
         return;
     }
 
-    auto replState = invariant(_getIndexBuild(buildUUID));
+    auto replState = invariant(_getIndexBuild(buildUUID),
+                               str::stream() << "Abort timestamp: " << abortTimestamp.toString());
 
     auto fut = replState->sharedPromise.getFuture();
     LOGV2(20655,
@@ -691,6 +711,67 @@ void IndexBuildsCoordinator::abortIndexBuildByBuildUUID(OperationContext* opCtx,
           "fut_waitNoThrow"_attr = fut.waitNoThrow());
 }
 
+boost::optional<UUID> IndexBuildsCoordinator::abortIndexBuildByIndexNamesNoWait(
+    OperationContext* opCtx,
+    const UUID& collectionUUID,
+    const std::vector<std::string>& indexNames,
+    Timestamp abortTimestamp,
+    const std::string& reason) {
+    boost::optional<UUID> buildUUID;
+    auto indexBuilds = _getIndexBuilds();
+    auto onIndexBuild = [&](std::shared_ptr<ReplIndexBuildState> replState) {
+        if (replState->collectionUUID != collectionUUID) {
+            return;
+        }
+
+        bool matchedBuilder = std::is_permutation(indexNames.begin(),
+                                                  indexNames.end(),
+                                                  replState->indexNames.begin(),
+                                                  replState->indexNames.end());
+        if (!matchedBuilder) {
+            return;
+        }
+
+        log() << "About to abort index builder: " << replState->buildUUID
+              << " on collection: " << collectionUUID
+              << ". First index: " << replState->indexNames.front();
+
+        if (this->abortIndexBuildByBuildUUIDNoWait(
+                opCtx, replState->buildUUID, abortTimestamp, reason)) {
+            buildUUID = replState->buildUUID;
+        }
+    };
+    forEachIndexBuild(indexBuilds,
+                      "IndexBuildsCoordinator::abortIndexBuildByIndexNamesNoWait - "_sd,
+                      onIndexBuild);
+    return buildUUID;
+}
+
+bool IndexBuildsCoordinator::hasIndexBuilder(OperationContext* opCtx,
+                                             const UUID& collectionUUID,
+                                             const std::vector<std::string>& indexNames) const {
+    bool foundIndexBuilder = false;
+    boost::optional<UUID> buildUUID;
+    auto indexBuilds = _getIndexBuilds();
+    auto onIndexBuild = [&](std::shared_ptr<ReplIndexBuildState> replState) {
+        if (replState->collectionUUID != collectionUUID) {
+            return;
+        }
+
+        bool matchedBuilder = std::is_permutation(indexNames.begin(),
+                                                  indexNames.end(),
+                                                  replState->indexNames.begin(),
+                                                  replState->indexNames.end());
+        if (!matchedBuilder) {
+            return;
+        }
+
+        foundIndexBuilder = true;
+    };
+    forEachIndexBuild(indexBuilds, "IndexBuildsCoordinator::hasIndexBuilder - "_sd, onIndexBuild);
+    return foundIndexBuilder;
+}
+
 bool IndexBuildsCoordinator::abortIndexBuildByBuildUUIDNoWait(OperationContext* opCtx,
                                                               const UUID& buildUUID,
                                                               Timestamp abortTimestamp,
@@ -913,6 +994,21 @@ void IndexBuildsCoordinator::assertNoBgOpInProgForDb(StringData db) const {
             !inProgForDb(db));
 }
 
+void IndexBuildsCoordinator::awaitIndexBuildFinished(const UUID& collectionUUID,
+                                                     const UUID& buildUUID) const {
+    stdx::unique_lock<Latch> lk(_mutex);
+
+    auto collIndexBuildsIt = _collectionIndexBuilds.find(collectionUUID);
+    if (collIndexBuildsIt == _collectionIndexBuilds.end()) {
+        return;
+    }
+
+    // Take a shared ptr, rather than accessing the Tracker through the map's iterator, so that the
+    // object does not destruct while we are waiting, causing a use-after-free memory error.
+    auto collIndexBuildsSharedPtr = collIndexBuildsIt->second;
+    collIndexBuildsSharedPtr->waitUntilIndexBuildFinished(lk, buildUUID);
+}
+
 void IndexBuildsCoordinator::awaitNoIndexBuildInProgressForCollection(
     const UUID& collectionUUID) const {
     stdx::unique_lock<Latch> lk(_mutex);
@@ -926,6 +1022,7 @@ void IndexBuildsCoordinator::awaitNoIndexBuildInProgressForCollection(
     // object does not destruct while we are waiting, causing a use-after-free memory error.
     auto collIndexBuildsSharedPtr = collIndexBuildsIt->second;
     collIndexBuildsSharedPtr->waitUntilNoIndexBuildsRemain(lk);
+    invariant(collIndexBuildsSharedPtr->getNumberOfIndexBuilds(lk) == 0);
 }
 
 void IndexBuildsCoordinator::awaitNoBgOpInProgForDb(StringData db) const {
@@ -1457,7 +1554,7 @@ void IndexBuildsCoordinator::_cleanUpSinglePhaseAfterFailure(
     if (status == ErrorCodes::InterruptedAtShutdown) {
         // Leave it as-if kill -9 happened. Startup recovery will rebuild the index.
         _indexBuildsManager.abortIndexBuildWithoutCleanup(
-            opCtx, replState->buildUUID, "shutting down");
+            opCtx, collection, replState->buildUUID, "shutting down");
         _indexBuildsManager.tearDownIndexBuild(
             opCtx, collection, replState->buildUUID, MultiIndexBlock::kNoopOnCleanUpFn);
         return;
@@ -1499,7 +1596,7 @@ void IndexBuildsCoordinator::_cleanUpTwoPhaseAfterFailure(
     if (status == ErrorCodes::InterruptedAtShutdown) {
         // Leave it as-if kill -9 happened. Startup recovery will restart the index build.
         _indexBuildsManager.abortIndexBuildWithoutCleanup(
-            opCtx, replState->buildUUID, "shutting down");
+            opCtx, collection, replState->buildUUID, "shutting down");
         _indexBuildsManager.tearDownIndexBuild(
             opCtx, collection, replState->buildUUID, MultiIndexBlock::kNoopOnCleanUpFn);
         return;
@@ -1532,7 +1629,7 @@ void IndexBuildsCoordinator::_cleanUpTwoPhaseAfterFailure(
             // rollback process to correct this state.
             if (abortIndexBuildTimestamp.isNull()) {
                 _indexBuildsManager.abortIndexBuildWithoutCleanup(
-                    opCtx, replState->buildUUID, "no longer primary");
+                    opCtx, collection, replState->buildUUID, "no longer primary");
                 _indexBuildsManager.tearDownIndexBuild(
                     opCtx, collection, replState->buildUUID, MultiIndexBlock::kNoopOnCleanUpFn);
                 return;
diff --git a/src/mongo/db/index_builds_coordinator.h b/src/mongo/db/index_builds_coordinator.h
index 0aa5c17a9b4..1be72d0554c 100644
--- a/src/mongo/db/index_builds_coordinator.h
+++ b/src/mongo/db/index_builds_coordinator.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <map>
 #include <string>
 #include <vector>
@@ -206,13 +207,19 @@ public:
      *     AutoGetCollection autoColl(..., collectionUUID, ...);
      *     autoColl->dropCollection(...);
      * }
-     *
-     * TODO: this is partially implemented. It calls IndexBuildsManager::abortIndexBuild that is not
-     * implemented.
      */
     void abortCollectionIndexBuilds(const UUID& collectionUUID, const std::string& reason);
 
     /**
+     * Signals all of the index builds on the specified collection to abort and returns the build
+     * UUIDs of the index builds that will be aborted. Must identify the collection with a UUID. The
+     * provided 'reason' will be used in the error message that the index builders return to their
+     * callers.
+     */
+    std::vector<UUID> abortCollectionIndexBuildsNoWait(const UUID& collectionUUID,
+                                                       const std::string& reason);
+
+    /**
      * Signals all of the index builds on the specified 'db' to abort and then waits until the index
      * builds are no longer running. The provided 'reason' will be used in the error message that
      * the index builders return to their callers.
@@ -227,9 +234,6 @@ public:
      *     AutoGetDb autoDb(...);
      *     autoDb->dropDatabase(...);
      * }
-     *
-     * TODO: this is partially implemented. It calls IndexBuildsManager::abortIndexBuild that is not
-     * implemented.
      */
     void abortDatabaseIndexBuilds(StringData db, const std::string& reason);
 
@@ -249,6 +253,27 @@ public:
                                           const UUID& buildUUID,
                                           Timestamp abortTimestamp,
                                           const std::string& reason);
+
+    /**
+     * Aborts an index build by its index name(s). This will only abort in-progress index builds if
+     * all of the indexes are specified that a single builder is building together. When an
+     * appropriate builder exists, this returns the build UUID of the index builder that will be
+     * aborted.
+     */
+    boost::optional<UUID> abortIndexBuildByIndexNamesNoWait(
+        OperationContext* opCtx,
+        const UUID& collectionUUID,
+        const std::vector<std::string>& indexNames,
+        Timestamp abortTimestamp,
+        const std::string& reason);
+
+    /**
+     * Returns true if there is an index builder building the given index names on a collection.
+     */
+    bool hasIndexBuilder(OperationContext* opCtx,
+                         const UUID& collectionUUID,
+                         const std::vector<std::string>& indexNames) const;
+
     /**
      * Returns number of index builds in process.
      *
@@ -327,6 +352,12 @@ public:
     void assertNoBgOpInProgForDb(StringData db) const;
 
     /**
+     * Waits for the index build with 'buildUUID' to finish on the specified collection before
+     * returning. Returns immediately if no such index build with 'buildUUID' is found.
+     */
+    void awaitIndexBuildFinished(const UUID& collectionUUID, const UUID& buildUUID) const;
+
+    /**
      * Waits for all index builds on a specified collection to finish.
      */
     void awaitNoIndexBuildInProgressForCollection(const UUID& collectionUUID) const;
@@ -633,6 +664,15 @@ protected:
      */
     std::vector<std::shared_ptr<ReplIndexBuildState>> _getIndexBuilds() const;
 
+    /**
+     * Helper for 'abortCollectionIndexBuilds' and 'abortCollectionIndexBuildsNoWait'. Returns the
+     * UUIDs of the aborted index builders
+     */
+    std::vector<UUID> _abortCollectionIndexBuilds(stdx::unique_lock<Latch>& lk,
+                                                  const UUID& collectionUUID,
+                                                  const std::string& reason,
+                                                  bool shouldWait);
+
     // Protects the below state.
     mutable Mutex _mutex = MONGO_MAKE_LATCH("IndexBuildsCoordinator::_mutex");