SERVER-38982 Fix RequiresIndexStage to avoid holding IndexCatalogEntry by shared_ptr.

2 files changed, 34 insertions, 33 deletions

diff --git a/src/mongo/db/exec/requires_index_stage.cpp b/src/mongo/db/exec/requires_index_stage.cpp
index 8718f98d1f4..307de4b7409 100644
--- a/src/mongo/db/exec/requires_index_stage.cpp
+++ b/src/mongo/db/exec/requires_index_stage.cpp
@@ -33,29 +33,42 @@
 
 namespace mongo {
 
+RequiresIndexStage::RequiresIndexStage(const char* stageType,
+                                       OperationContext* opCtx,
+                                       const IndexDescriptor* indexDescriptor)
+    : RequiresCollectionStage(stageType, opCtx, indexDescriptor->getCollection()),
+      _weakIndexCatalogEntry(collection()->getIndexCatalog()->getEntryShared(indexDescriptor)) {
+    auto indexCatalogEntry = _weakIndexCatalogEntry.lock();
+    _indexDescriptor = indexCatalogEntry->descriptor();
+    _indexAccessMethod = indexCatalogEntry->accessMethod();
+    invariant(_indexDescriptor);
+    invariant(_indexAccessMethod);
+    _indexName = _indexDescriptor->indexName();
+}
+
 void RequiresIndexStage::doSaveStateRequiresCollection() {
     doSaveStateRequiresIndex();
 
-    // During yield, we relinquish our shared ownership of the index catalog entry. This allows the
-    // index to be dropped during yield, but permits us to check via the weak_ptr interface
-    // whether the index is still valid on yield recovery.
-    //
-    // We also set catalog pointers to null, since accessing these pointers is illegal during yield.
-    _indexCatalogEntry.reset();
+    // Set catalog pointers to null, since accessing these pointers is illegal during yield.
     _indexDescriptor = nullptr;
     _indexAccessMethod = nullptr;
 }
 
 void RequiresIndexStage::doRestoreStateRequiresCollection() {
-    // Reacquire shared ownership of the index catalog entry. If we're unable to do so, then the
-    // our index is no longer valid, and the query should die.
-    _indexCatalogEntry = _weakIndexCatalogEntry.lock();
+    // Attempt to lock the weak_ptr. If the resulting shared_ptr is null, then our index is no
+    // longer valid and the query should die.
+    auto indexCatalogEntry = _weakIndexCatalogEntry.lock();
     uassert(ErrorCodes::QueryPlanKilled,
             str::stream() << "query plan killed :: index '" << _indexName << "' dropped",
-            _indexCatalogEntry);
+            indexCatalogEntry);
 
-    _indexDescriptor = _indexCatalogEntry->descriptor();
-    _indexAccessMethod = _indexCatalogEntry->accessMethod();
+    // Re-obtain catalog pointers that were set to null during yield preparation. It is safe to
+    // access the catalog entry by raw pointer when the query is active, as its validity is
+    // protected by at least MODE_IS collection locks.
+    _indexDescriptor = indexCatalogEntry->descriptor();
+    _indexAccessMethod = indexCatalogEntry->accessMethod();
+    invariant(_indexDescriptor);
+    invariant(_indexAccessMethod);
 
     doRestoreStateRequiresIndex();
 }
diff --git a/src/mongo/db/exec/requires_index_stage.h b/src/mongo/db/exec/requires_index_stage.h
index e941e20ae4d..030aa369306 100644
--- a/src/mongo/db/exec/requires_index_stage.h
+++ b/src/mongo/db/exec/requires_index_stage.h
@@ -48,17 +48,7 @@ class RequiresIndexStage : public RequiresCollectionStage {
 public:
     RequiresIndexStage(const char* stageType,
                        OperationContext* opCtx,
-                       const IndexDescriptor* indexDescriptor)
-        : RequiresCollectionStage(stageType, opCtx, indexDescriptor->getCollection()),
-          _weakIndexCatalogEntry(collection()->getIndexCatalog()->getEntryShared(indexDescriptor)),
-          _indexCatalogEntry(_weakIndexCatalogEntry.lock()),
-          _indexDescriptor(indexDescriptor),
-          _indexAccessMethod(_indexCatalogEntry->accessMethod()),
-          _indexName(_indexDescriptor->indexName()) {
-        invariant(_indexCatalogEntry);
-        invariant(_indexDescriptor);
-        invariant(_indexAccessMethod);
-    }
+                       const IndexDescriptor* indexDescriptor);
 
     virtual ~RequiresIndexStage() = default;
 
@@ -86,25 +76,23 @@ protected:
     }
 
 private:
-    // This stage shares ownership of the index catalog entry when the query is running, and
-    // relinquishes its shared ownership when in a saved state for a yield or between getMores. We
-    // keep a weak_ptr to the entry in order to reacquire shared ownership on yield recovery,
-    // throwing a query-fatal exception if the weak_ptr indicates that the underlying catalog object
-    // has been destroyed.
+    // We keep a weak_ptr to the index catalog entry in order to detect when the underlying catalog
+    // object has been destroyed, e.g. due to an index drop. In this scenario, the
+    // RequiresIndexStage will throw a query-fatal exception after an attempt to lock the weak_ptr
+    // indicates that the pointed-to object was deleted.
     //
     // This is necessary to protect against that case that our index is dropped and then recreated
     // during yield. Such an event should cause the query to be killed, since index cursors may have
     // pointers into catalog objects that no longer exist. Since indices do not have UUIDs,
-    // different epochs of the index cannot be distinguished. The weak_ptr allows us to relinquish
-    // ownership of the index during yield, but also determine whether the pointed-to object has
-    // been destroyed during yield recovery.
+    // different epochs of the index cannot be distinguished. The weak_ptr allows us to safely.
+    // determine whether the pointed-to object has been destroyed during yield recovery via the
+    // shared_ptr control block.
     std::weak_ptr<const IndexCatalogEntry> _weakIndexCatalogEntry;
-    std::shared_ptr<const IndexCatalogEntry> _indexCatalogEntry;
 
     const IndexDescriptor* _indexDescriptor;
     const IndexAccessMethod* _indexAccessMethod;
 
-    const std::string _indexName;
+    std::string _indexName;
 };
 
 }  // namespace mongo