SERVER-39703 Pull opCtx ptr out of the TemporaryKVRecordStore and instead pass it as a function parameter

17 files changed, 102 insertions, 23 deletions

diff --git a/src/mongo/db/catalog/database_test.cpp b/src/mongo/db/catalog/database_test.cpp
index 230e2f23251..9b44575217e 100644
--- a/src/mongo/db/catalog/database_test.cpp
+++ b/src/mongo/db/catalog/database_test.cpp
@@ -318,6 +318,7 @@ void _testDropCollectionThrowsExceptionIfThereAreIndexesInProgress(OperationCont
             WriteUnitOfWork wuow(opCtx);
             indexBuildBlock->success(opCtx, collection);
             wuow.commit();
+            indexBuildBlock->deleteTemporaryTables(opCtx);
         });
 
         ASSERT_GREATER_THAN(indexCatalog->numIndexesInProgress(opCtx), 0);
diff --git a/src/mongo/db/catalog/index_build_block.cpp b/src/mongo/db/catalog/index_build_block.cpp
index 3a3fb33eff1..1715fbf1c50 100644
--- a/src/mongo/db/catalog/index_build_block.cpp
+++ b/src/mongo/db/catalog/index_build_block.cpp
@@ -52,6 +52,12 @@ IndexCatalogImpl::IndexBuildBlock::IndexBuildBlock(IndexCatalogImpl* catalog,
                                                    IndexBuildMethod method)
     : _catalog(catalog), _ns(nss.ns()), _spec(spec.getOwned()), _method(method), _entry(nullptr) {}
 
+void IndexCatalogImpl::IndexBuildBlock::deleteTemporaryTables(OperationContext* opCtx) {
+    if (_indexBuildInterceptor) {
+        _indexBuildInterceptor->deleteTemporaryTables(opCtx);
+    }
+}
+
 Status IndexCatalogImpl::IndexBuildBlock::init(OperationContext* opCtx, Collection* collection) {
     // Being in a WUOW means all timestamping responsibility can be pushed up to the caller.
     invariant(opCtx->lockState()->inAWriteUnitOfWork());
diff --git a/src/mongo/db/catalog/index_catalog.h b/src/mongo/db/catalog/index_catalog.h
index 59b4cb249cf..5ee4aed313c 100644
--- a/src/mongo/db/catalog/index_catalog.h
+++ b/src/mongo/db/catalog/index_catalog.h
@@ -156,6 +156,12 @@ public:
         virtual ~IndexBuildBlockInterface() = default;
 
         /**
+         * Must be called before the object is destructed if init() has been called.
+         * Cleans up the temporary tables that are created for an index build.
+         */
+        virtual void deleteTemporaryTables(OperationContext* opCtx) = 0;
+
+        /**
          * Initializes a new entry for the index in the IndexCatalog.
          *
          * On success, holds pointer to newly created IndexCatalogEntry that can be accessed using
diff --git a/src/mongo/db/catalog/index_catalog_impl.h b/src/mongo/db/catalog/index_catalog_impl.h
index 943395c1df6..933bab58b62 100644
--- a/src/mongo/db/catalog/index_catalog_impl.h
+++ b/src/mongo/db/catalog/index_catalog_impl.h
@@ -252,6 +252,11 @@ public:
         ~IndexBuildBlock();
 
         /**
+         * Being called in a 'WriteUnitOfWork' has no effect.
+         */
+        void deleteTemporaryTables(OperationContext* opCtx);
+
+        /**
          * Must be called from within a `WriteUnitOfWork`
          */
         Status init(OperationContext* opCtx, Collection* collection);
diff --git a/src/mongo/db/catalog/multi_index_block.cpp b/src/mongo/db/catalog/multi_index_block.cpp
index 552f0eafc49..55f8bf7bd80 100644
--- a/src/mongo/db/catalog/multi_index_block.cpp
+++ b/src/mongo/db/catalog/multi_index_block.cpp
@@ -95,6 +95,12 @@ void MultiIndexBlock::cleanUpAfterBuild(OperationContext* opCtx, Collection* col
     }
 
     if (!_needToCleanup || _indexes.empty()) {
+        // The temp tables cannot be dropped in commit() because commit() can be called multiple
+        // times on write conflict errors and the drop does not rollback in WUOWs.
+        for (auto& index : _indexes) {
+            index.block->deleteTemporaryTables(opCtx);
+        }
+
         _buildIsCleanedUp = true;
         return;
     }
@@ -109,6 +115,7 @@ void MultiIndexBlock::cleanUpAfterBuild(OperationContext* opCtx, Collection* col
             // a WUOW. Nothing inside this block can fail, and it is made fatal if it does.
             for (size_t i = 0; i < _indexes.size(); i++) {
                 _indexes[i].block->fail(opCtx, collection);
+                _indexes[i].block->deleteTemporaryTables(opCtx);
             }
 
             auto replCoord = repl::ReplicationCoordinator::get(opCtx);
@@ -228,7 +235,12 @@ StatusWith<std::vector<BSONObj>> MultiIndexBlock::init(OperationContext* opCtx,
     // On rollback in init(), cleans up _indexes so that ~MultiIndexBlock doesn't try to clean up
     // _indexes manually (since the changes were already rolled back).
     // Due to this, it is thus legal to call init() again after it fails.
-    opCtx->recoveryUnit()->onRollback([this]() { _indexes.clear(); });
+    opCtx->recoveryUnit()->onRollback([this, opCtx]() {
+        for (auto& index : _indexes) {
+            index.block->deleteTemporaryTables(opCtx);
+        }
+        _indexes.clear();
+    });
 
     const auto& ns = collection->ns().ns();
 
@@ -630,7 +642,6 @@ Status MultiIndexBlock::drainBackgroundWrites(OperationContext* opCtx,
     return Status::OK();
 }
 
-
 Status MultiIndexBlock::checkConstraints(OperationContext* opCtx) {
     if (State::kAborted == _getState()) {
         return {ErrorCodes::IndexBuildAborted,
@@ -657,8 +668,11 @@ Status MultiIndexBlock::checkConstraints(OperationContext* opCtx) {
     return Status::OK();
 }
 
-void MultiIndexBlock::abortWithoutCleanup() {
+void MultiIndexBlock::abortWithoutCleanup(OperationContext* opCtx) {
     _setStateToAbortedIfNotCommitted("aborted without cleanup"_sd);
+    for (auto& index : _indexes) {
+        index.block->deleteTemporaryTables(opCtx);
+    }
     _indexes.clear();
     _needToCleanup = false;
 }
diff --git a/src/mongo/db/catalog/multi_index_block.h b/src/mongo/db/catalog/multi_index_block.h
index 98368293411..9ef56a6a998 100644
--- a/src/mongo/db/catalog/multi_index_block.h
+++ b/src/mongo/db/catalog/multi_index_block.h
@@ -260,7 +260,7 @@ public:
      *
      * Must be called from owning thread.
      */
-    void abortWithoutCleanup();
+    void abortWithoutCleanup(OperationContext* opCtx);
 
     /**
      * Returns true if this build block supports background writes while building an index. This is
diff --git a/src/mongo/db/catalog/multi_index_block_test.cpp b/src/mongo/db/catalog/multi_index_block_test.cpp
index e1e8c30133b..39adbc0728d 100644
--- a/src/mongo/db/catalog/multi_index_block_test.cpp
+++ b/src/mongo/db/catalog/multi_index_block_test.cpp
@@ -165,7 +165,7 @@ TEST_F(MultiIndexBlockTest, AbortWithoutCleanupAfterInsertingSingleDocument) {
         getOpCtx(), getCollection(), std::vector<BSONObj>(), MultiIndexBlock::kNoopOnInitFn));
     ASSERT_EQUALS(0U, specs.size());
     ASSERT_OK(indexer->insert(getOpCtx(), {}, {}));
-    indexer->abortWithoutCleanup();
+    indexer->abortWithoutCleanup(getOpCtx());
     ASSERT_EQUALS(MultiIndexBlock::State::kAborted, indexer->getState_forTest());
 
     ASSERT_FALSE(indexer->isCommitted());
diff --git a/src/mongo/db/index/duplicate_key_tracker.cpp b/src/mongo/db/index/duplicate_key_tracker.cpp
index 55e93838206..ac06389c2d3 100644
--- a/src/mongo/db/index/duplicate_key_tracker.cpp
+++ b/src/mongo/db/index/duplicate_key_tracker.cpp
@@ -54,6 +54,10 @@ DuplicateKeyTracker::DuplicateKeyTracker(OperationContext* opCtx, const IndexCat
     invariant(_indexCatalogEntry->descriptor()->unique());
 }
 
+void DuplicateKeyTracker::deleteTemporaryTable(OperationContext* opCtx) {
+    _keyConstraintsTable->deleteTemporaryTable(opCtx);
+}
+
 Status DuplicateKeyTracker::recordKeys(OperationContext* opCtx, const std::vector<BSONObj>& keys) {
     if (keys.size() == 0)
         return Status::OK();
diff --git a/src/mongo/db/index/duplicate_key_tracker.h b/src/mongo/db/index/duplicate_key_tracker.h
index 8246397a508..3879cdb3486 100644
--- a/src/mongo/db/index/duplicate_key_tracker.h
+++ b/src/mongo/db/index/duplicate_key_tracker.h
@@ -50,9 +50,19 @@ class DuplicateKeyTracker {
     MONGO_DISALLOW_COPYING(DuplicateKeyTracker);
 
 public:
+    /**
+     * Creates a temporary table in which to store any duplicate key constraint violations.
+     * deleteTemporaryTable() must be called before destruction.
+     */
     DuplicateKeyTracker(OperationContext* opCtx, const IndexCatalogEntry* indexCatalogEntry);
 
     /**
+     * Deletes the temporary table for the duplicate key constraint violations. Must be called
+     * before object destruction.
+     */
+    void deleteTemporaryTable(OperationContext* opCtx);
+
+    /**
      * Given a set of duplicate keys, insert them into the key constraint table.
      */
     Status recordKeys(OperationContext* opCtx, const std::vector<BSONObj>& keys);
diff --git a/src/mongo/db/index/index_build_interceptor.cpp b/src/mongo/db/index/index_build_interceptor.cpp
index 8b8ae089912..416a9203746 100644
--- a/src/mongo/db/index/index_build_interceptor.cpp
+++ b/src/mongo/db/index/index_build_interceptor.cpp
@@ -62,6 +62,13 @@ IndexBuildInterceptor::IndexBuildInterceptor(OperationContext* opCtx, IndexCatal
     }
 }
 
+void IndexBuildInterceptor::deleteTemporaryTables(OperationContext* opCtx) {
+    _sideWritesTable->deleteTemporaryTable(opCtx);
+    if (_duplicateKeyTracker) {
+        _duplicateKeyTracker->deleteTemporaryTable(opCtx);
+    }
+}
+
 Status IndexBuildInterceptor::recordDuplicateKeys(OperationContext* opCtx,
                                                   const std::vector<BSONObj>& keys) {
     invariant(_indexCatalogEntry->descriptor()->unique());
diff --git a/src/mongo/db/index/index_build_interceptor.h b/src/mongo/db/index/index_build_interceptor.h
index e4f2f2b3027..ccaa261e425 100644
--- a/src/mongo/db/index/index_build_interceptor.h
+++ b/src/mongo/db/index/index_build_interceptor.h
@@ -47,13 +47,21 @@ public:
     enum class Op { kInsert, kDelete };
 
     /**
-     * The OperationContext is used to construct a temporary table in the storage engine to
-     * intercept side writes. This interceptor must not exist longer than the operation context used
-     * to construct it, as the underlying TemporaryRecordStore needs it to destroy itself.
+     * Creates a temporary table for writes during an index build. Additionally creates a temporary
+     * table to store any duplicate key constraint violations found during the build, if the index
+     * being built has uniqueness constraints.
+     *
+     * deleteTemporaryTable() must be called before destruction to delete the temporary tables.
      */
     IndexBuildInterceptor(OperationContext* opCtx, IndexCatalogEntry* entry);
 
     /**
+     * Deletes the temporary side writes and duplicate key constraint violations tables. Must be
+     * called before object destruction.
+     */
+    void deleteTemporaryTables(OperationContext* opCtx);
+
+    /**
      * Client writes that are concurrent with an index build will have their index updates written
      * to a temporary table. After the index table scan is complete, these updates will be applied
      * to the underlying index table.
diff --git a/src/mongo/db/index_builder.cpp b/src/mongo/db/index_builder.cpp
index a40d5bfc24a..4d9c82548b0 100644
--- a/src/mongo/db/index_builder.cpp
+++ b/src/mongo/db/index_builder.cpp
@@ -185,7 +185,7 @@ Status IndexBuilder::_buildAndHandleErrors(OperationContext* opCtx,
 
     if (status.code() == ErrorCodes::InterruptedAtShutdown) {
         // leave it as-if kill -9 happened. This will be handled on restart.
-        indexer.abortWithoutCleanup();
+        indexer.abortWithoutCleanup(opCtx);
         return status;
     }
 
diff --git a/src/mongo/db/storage/kv/kv_storage_engine.cpp b/src/mongo/db/storage/kv/kv_storage_engine.cpp
index ca9f61e6431..b10de937bb7 100644
--- a/src/mongo/db/storage/kv/kv_storage_engine.cpp
+++ b/src/mongo/db/storage/kv/kv_storage_engine.cpp
@@ -709,7 +709,7 @@ std::unique_ptr<TemporaryRecordStore> KVStorageEngine::makeTemporaryRecordStore(
     std::unique_ptr<RecordStore> rs =
         _engine->makeTemporaryRecordStore(opCtx, _catalog->newInternalIdent());
     LOG(1) << "created temporary record store: " << rs->getIdent();
-    return std::make_unique<TemporaryKVRecordStore>(opCtx, getEngine(), std::move(rs));
+    return std::make_unique<TemporaryKVRecordStore>(getEngine(), std::move(rs));
 }
 
 void KVStorageEngine::setJournalListener(JournalListener* jl) {
diff --git a/src/mongo/db/storage/kv/kv_storage_engine_test.cpp b/src/mongo/db/storage/kv/kv_storage_engine_test.cpp
index 45e203d8b01..9f5d748e7f9 100644
--- a/src/mongo/db/storage/kv/kv_storage_engine_test.cpp
+++ b/src/mongo/db/storage/kv/kv_storage_engine_test.cpp
@@ -277,6 +277,8 @@ TEST_F(KVStorageEngineTest, ReconcileDropsTemporary) {
 
     // The storage engine is responsible for dropping its temporary idents.
     ASSERT(!identExists(opCtx.get(), ident));
+
+    rs->deleteTemporaryTable(opCtx.get());
 }
 
 TEST_F(KVStorageEngineTest, TemporaryDropsItself) {
@@ -289,6 +291,8 @@ TEST_F(KVStorageEngineTest, TemporaryDropsItself) {
         ident = rs->rs()->getIdent();
 
         ASSERT(identExists(opCtx.get(), ident));
+
+        rs->deleteTemporaryTable(opCtx.get());
     }
 
     // The temporary record store RAII class should drop itself.
@@ -330,6 +334,9 @@ TEST_F(KVStorageEngineTest, ReconcileDoesNotDropIndexBuildTempTables) {
     // The owning index was dropped, and so should its temporary tables.
     ASSERT(!identExists(opCtx.get(), sideWrites->rs()->getIdent()));
     ASSERT(!identExists(opCtx.get(), constraintViolations->rs()->getIdent()));
+
+    sideWrites->deleteTemporaryTable(opCtx.get());
+    constraintViolations->deleteTemporaryTable(opCtx.get());
 }
 
 TEST_F(KVStorageEngineTest, ReconcileDoesNotDropIndexBuildTempTablesBackgroundSecondary) {
@@ -372,6 +379,9 @@ TEST_F(KVStorageEngineTest, ReconcileDoesNotDropIndexBuildTempTablesBackgroundSe
     // dropped.
     ASSERT(identExists(opCtx.get(), sideWrites->rs()->getIdent()));
     ASSERT(identExists(opCtx.get(), constraintViolations->rs()->getIdent()));
+
+    sideWrites->deleteTemporaryTable(opCtx.get());
+    constraintViolations->deleteTemporaryTable(opCtx.get());
 }
 
 TEST_F(KVStorageEngineRepairTest, LoadCatalogRecoversOrphans) {
diff --git a/src/mongo/db/storage/kv/temporary_kv_record_store.cpp b/src/mongo/db/storage/kv/temporary_kv_record_store.cpp
index dffcb864268..c752e9fca05 100644
--- a/src/mongo/db/storage/kv/temporary_kv_record_store.cpp
+++ b/src/mongo/db/storage/kv/temporary_kv_record_store.cpp
@@ -40,10 +40,15 @@
 namespace mongo {
 
 TemporaryKVRecordStore::~TemporaryKVRecordStore() {
-    auto status = _kvEngine->dropIdent(_opCtx, _rs->getIdent());
+    invariant(_recordStoreHasBeenDeleted);
+}
+
+void TemporaryKVRecordStore::deleteTemporaryTable(OperationContext* opCtx) {
+    auto status = _kvEngine->dropIdent(opCtx, _rs->getIdent());
     fassert(
         51032,
         status.withContext(str::stream() << "failed to drop temporary ident: " << _rs->getIdent()));
+    _recordStoreHasBeenDeleted = true;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/kv/temporary_kv_record_store.h b/src/mongo/db/storage/kv/temporary_kv_record_store.h
index c9b7a6945ad..a992ca69cd4 100644
--- a/src/mongo/db/storage/kv/temporary_kv_record_store.h
+++ b/src/mongo/db/storage/kv/temporary_kv_record_store.h
@@ -38,17 +38,14 @@ class KVEngine;
 class OperationContext;
 
 /**
- * This is an implementation of an RAII type that manages a temporary RecordStore on a KVEngine.
+ * Implementation of TemporaryRecordStore that manages a temporary RecordStore on a KVEngine.
  *
- * This object should not exist any longer than the provided OperationContext, as the destructor
- * uses it to drop the record store on the KVEngine.
+ * deleteTemporaryTable() must be called before destruction to delete the underlying RecordStore.
  */
 class TemporaryKVRecordStore : public TemporaryRecordStore {
 public:
-    TemporaryKVRecordStore(OperationContext* opCtx,
-                           KVEngine* kvEngine,
-                           std::unique_ptr<RecordStore> rs)
-        : TemporaryRecordStore(std::move(rs)), _opCtx(opCtx), _kvEngine(kvEngine){};
+    TemporaryKVRecordStore(KVEngine* kvEngine, std::unique_ptr<RecordStore> rs)
+        : TemporaryRecordStore(std::move(rs)), _kvEngine(kvEngine){};
 
     // Not copyable.
     TemporaryKVRecordStore(const TemporaryKVRecordStore&) = delete;
@@ -57,14 +54,18 @@ public:
     // Move constructor.
     TemporaryKVRecordStore(TemporaryKVRecordStore&& other) noexcept
         : TemporaryRecordStore(std::move(other._rs)),
-          _opCtx(other._opCtx),
           _kvEngine(other._kvEngine) {}
 
     ~TemporaryKVRecordStore();
 
+    /**
+     * Drops the persisted record store from the storage engine.
+     */
+    void deleteTemporaryTable(OperationContext* opCtx);
+
 private:
-    OperationContext* _opCtx;
     KVEngine* _kvEngine;
+    bool _recordStoreHasBeenDeleted = false;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/temporary_record_store.h b/src/mongo/db/storage/temporary_record_store.h
index b42f6aafad0..21859790415 100644
--- a/src/mongo/db/storage/temporary_record_store.h
+++ b/src/mongo/db/storage/temporary_record_store.h
@@ -34,10 +34,10 @@
 namespace mongo {
 
 /**
- * This is an RAII type that manages the lifetime of a temporary RecordStore.
+ * Manages the lifetime of a temporary RecordStore.
  *
- * Derived classes must implement a destructor that drops the underlying RecordStore from the
- * storage engine.
+ * Derived classes must implement and call deleteTemporaryTable() to delete the underlying
+ * RecordStore from the storage engine.
  */
 class TemporaryRecordStore {
 public:
@@ -52,6 +52,8 @@ public:
 
     virtual ~TemporaryRecordStore() {}
 
+    virtual void deleteTemporaryTable(OperationContext* opCtx) {}
+
     RecordStore* rs() {
         return _rs.get();
     }