SERVER-21057: Undo MMAPv1 invalidations on rollback

17 files changed, 50 insertions, 19 deletions

diff --git a/src/mongo/db/exec/collection_scan.cpp b/src/mongo/db/exec/collection_scan.cpp
index 613e5e14519..74d64c39cc8 100644
--- a/src/mongo/db/exec/collection_scan.cpp
+++ b/src/mongo/db/exec/collection_scan.cpp
@@ -203,7 +203,7 @@ void CollectionScan::doInvalidate(OperationContext* txn,
 
     // Deletions can harm the underlying RecordCursor so we must pass them down.
     if (_cursor) {
-        _cursor->invalidate(id);
+        _cursor->invalidate(txn, id);
     }
 
     if (_params.tailable && id == _lastSeenId) {
diff --git a/src/mongo/db/exec/multi_iterator.cpp b/src/mongo/db/exec/multi_iterator.cpp
index 8b881c87cf2..ffa3ece0b4e 100644
--- a/src/mongo/db/exec/multi_iterator.cpp
+++ b/src/mongo/db/exec/multi_iterator.cpp
@@ -137,7 +137,7 @@ void MultiIteratorStage::doInvalidate(OperationContext* txn,
     switch (type) {
         case INVALIDATION_DELETION:
             for (size_t i = 0; i < _iterators.size(); i++) {
-                _iterators[i]->invalidate(dl);
+                _iterators[i]->invalidate(txn, dl);
             }
             break;
         case INVALIDATION_MUTATION:
diff --git a/src/mongo/db/exec/oplogstart.cpp b/src/mongo/db/exec/oplogstart.cpp
index 9554f5f68e2..3bf2c17fe5b 100644
--- a/src/mongo/db/exec/oplogstart.cpp
+++ b/src/mongo/db/exec/oplogstart.cpp
@@ -175,7 +175,7 @@ void OplogStart::doInvalidate(OperationContext* txn, const RecordId& dl, Invalid
     }
 
     for (size_t i = 0; i < _subIterators.size(); i++) {
-        _subIterators[i]->invalidate(dl);
+        _subIterators[i]->invalidate(txn, dl);
     }
 }
 
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_base.cpp b/src/mongo/db/storage/mmap_v1/record_store_v1_base.cpp
index 9a7a841aca5..5f5e0438d73 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_base.cpp
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_base.cpp
@@ -919,8 +919,13 @@ void RecordStoreV1Base::IntraExtentIterator::advance() {
     _curr = (nextOfs == DiskLoc::NullOfs ? DiskLoc() : DiskLoc(_curr.a(), nextOfs));
 }
 
-void RecordStoreV1Base::IntraExtentIterator::invalidate(const RecordId& rid) {
+void RecordStoreV1Base::IntraExtentIterator::invalidate(OperationContext* txn,
+                                                        const RecordId& rid) {
     if (rid == _curr.toRecordId()) {
+        const DiskLoc origLoc = _curr;
+
+        // Undo the advance on rollback, as the deletion that forced it "never happened".
+        txn->recoveryUnit()->onRollback([this, origLoc]() { this->_curr = origLoc; });
         advance();
     }
 }
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_base.h b/src/mongo/db/storage/mmap_v1/record_store_v1_base.h
index 0489adced45..d34c5a9b3f0 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_base.h
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_base.h
@@ -333,7 +333,7 @@ public:
         : _txn(txn), _curr(start), _rs(rs), _forward(forward) {}
 
     boost::optional<Record> next() final;
-    void invalidate(const RecordId& dl) final;
+    void invalidate(OperationContext* txn, const RecordId& dl) final;
     void save() final {}
     bool restore() final {
         return true;
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.cpp b/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.cpp
index 96c50e73024..cdd8363d949 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.cpp
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.cpp
@@ -86,7 +86,7 @@ boost::optional<Record> CappedRecordStoreV1Iterator::seekExact(const RecordId& i
     return {{id, _recordStore->RecordStore::dataFor(_txn, id)}};
 }
 
-void CappedRecordStoreV1Iterator::invalidate(const RecordId& id) {
+void CappedRecordStoreV1Iterator::invalidate(OperationContext* txn, const RecordId& id) {
     const DiskLoc dl = DiskLoc::fromRecordId(id);
     if (dl == _curr) {
         // We *could* move to the next thing, since there is actually a next
@@ -94,6 +94,8 @@ void CappedRecordStoreV1Iterator::invalidate(const RecordId& id) {
         // "note we cannot advance here. if this condition occurs, writes to the oplog
         //  have "caught" the reader.  skipping ahead, the reader would miss potentially
         //  important data."
+        // We don't really need to worry about rollback here, as the very next write would
+        // invalidate the cursor anyway.
         _curr = DiskLoc();
         _killedByInvalidate = true;
     }
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.h b/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.h
index fbc47c01512..275c78cae38 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.h
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_capped_iterator.h
@@ -57,7 +57,7 @@ public:
     void reattachToOperationContext(OperationContext* txn) final {
         _txn = txn;
     }
-    void invalidate(const RecordId& dl) final;
+    void invalidate(OperationContext* txn, const RecordId& dl) final;
     std::unique_ptr<RecordFetcher> fetcherForNext() const final;
     std::unique_ptr<RecordFetcher> fetcherForId(const RecordId& id) const final;
 
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.cpp b/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.cpp
index 9d1d5a076ef..ac8f083eb82 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.cpp
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.cpp
@@ -181,7 +181,7 @@ bool RecordStoreV1RepairCursor::_advanceToNextValidExtent() {
     return true;
 }
 
-void RecordStoreV1RepairCursor::invalidate(const RecordId& id) {
+void RecordStoreV1RepairCursor::invalidate(OperationContext* txn, const RecordId& id) {
     // If we see this record again it probably means it was reinserted rather than an infinite
     // loop. If we do loop, we should quickly hit another seen record that hasn't been
     // invalidated.
@@ -191,6 +191,8 @@ void RecordStoreV1RepairCursor::invalidate(const RecordId& id) {
     if (_currRecord == dl) {
         // The DiskLoc being invalidated is also the one pointed at by this iterator. We
         // advance the iterator so it's not pointing at invalid data.
+        // We don't worry about undoing invalidations on rollback here, as we shouldn't have
+        // concurrent writes that can rollback to a database we're trying to recover.
         advance();
 
         if (_currRecord == dl) {
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.h b/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.h
index 5ece66ae39c..a45cb1ca9e7 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.h
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_repair_iterator.h
@@ -45,7 +45,7 @@ public:
     RecordStoreV1RepairCursor(OperationContext* txn, const RecordStoreV1Base* recordStore);
 
     boost::optional<Record> next() final;
-    void invalidate(const RecordId& dl);
+    void invalidate(OperationContext* txn, const RecordId& dl);
     void save() final {}
     bool restore() final {
         return true;
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.cpp b/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.cpp
index eb4f0ad6745..81cd3456a07 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.cpp
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.cpp
@@ -101,9 +101,13 @@ void SimpleRecordStoreV1Iterator::advance() {
     }
 }
 
-void SimpleRecordStoreV1Iterator::invalidate(const RecordId& dl) {
+void SimpleRecordStoreV1Iterator::invalidate(OperationContext* txn, const RecordId& dl) {
     // Just move past the thing being deleted.
     if (dl == _curr.toRecordId()) {
+        const DiskLoc origLoc = _curr;
+
+        // Undo the advance on rollback, as the deletion that forced it "never happened".
+        txn->recoveryUnit()->onRollback([this, origLoc]() { this->_curr = origLoc; });
         advance();
     }
 }
diff --git a/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.h b/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.h
index 4b3cea6cd6b..a480566f9d7 100644
--- a/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.h
+++ b/src/mongo/db/storage/mmap_v1/record_store_v1_simple_iterator.h
@@ -57,7 +57,7 @@ public:
     void reattachToOperationContext(OperationContext* txn) final {
         _txn = txn;
     }
-    void invalidate(const RecordId& dl) final;
+    void invalidate(OperationContext* txn, const RecordId& dl) final;
     std::unique_ptr<RecordFetcher> fetcherForNext() const final;
     std::unique_ptr<RecordFetcher> fetcherForId(const RecordId& id) const final;
 
diff --git a/src/mongo/db/storage/record_store.h b/src/mongo/db/storage/record_store.h
index 2ed43d9832c..b6e973a16ab 100644
--- a/src/mongo/db/storage/record_store.h
+++ b/src/mongo/db/storage/record_store.h
@@ -194,13 +194,13 @@ public:
     virtual void reattachToOperationContext(OperationContext* opCtx) = 0;
 
     /**
-     * Inform the cursor that this id is being invalidated.
-     * Must be called between save and restore.
+     * Inform the cursor that this id is being invalidated. Must be called between save and restore.
+     * The txn is that of the operation causing the invalidation, not the txn using the cursor.
      *
      * WARNING: Storage engines other than MMAPv1 should use the default implementation,
      *          and not depend on this being called.
      */
-    virtual void invalidate(const RecordId& id){};
+    virtual void invalidate(OperationContext* txn, const RecordId& id) {}
 
     //
     // RecordFetchers
diff --git a/src/mongo/db/storage/record_store_test_repairiter.cpp b/src/mongo/db/storage/record_store_test_repairiter.cpp
index 48dae6bcf1c..fd71f1a7523 100644
--- a/src/mongo/db/storage/record_store_test_repairiter.cpp
+++ b/src/mongo/db/storage/record_store_test_repairiter.cpp
@@ -157,7 +157,7 @@ TEST(RecordStoreTestHarness, GetIteratorForRepairInvalidateSingleton) {
 
         // Invalidate the record we're pointing at.
         cursor->save();
-        cursor->invalidate(idToInvalidate);
+        cursor->invalidate(opCtx.get(), idToInvalidate);
         cursor->restore();
 
         // Iterator should be EOF now because the only thing in the collection got deleted.
diff --git a/src/mongo/dbtests/query_stage_collscan.cpp b/src/mongo/dbtests/query_stage_collscan.cpp
index eea912d4bd4..d4307f54fbf 100644
--- a/src/mongo/dbtests/query_stage_collscan.cpp
+++ b/src/mongo/dbtests/query_stage_collscan.cpp
@@ -299,7 +299,11 @@ public:
 
         // Remove locs[count].
         scan->saveState();
-        scan->invalidate(&_txn, locs[count], INVALIDATION_DELETION);
+        {
+            WriteUnitOfWork wunit(&_txn);
+            scan->invalidate(&_txn, locs[count], INVALIDATION_DELETION);
+            wunit.commit();  // to avoid rollback of the invalidate
+        }
         remove(coll->docFor(&_txn, locs[count]).value());
         scan->restoreState();
 
@@ -360,7 +364,11 @@ public:
 
         // Remove locs[count].
         scan->saveState();
-        scan->invalidate(&_txn, locs[count], INVALIDATION_DELETION);
+        {
+            WriteUnitOfWork wunit(&_txn);
+            scan->invalidate(&_txn, locs[count], INVALIDATION_DELETION);
+            wunit.commit();  // to avoid rollback of the invalidate
+        }
         remove(coll->docFor(&_txn, locs[count]).value());
         scan->restoreState();
 
diff --git a/src/mongo/dbtests/query_stage_count.cpp b/src/mongo/dbtests/query_stage_count.cpp
index b1273f6f7f8..9065c880ff6 100644
--- a/src/mongo/dbtests/query_stage_count.cpp
+++ b/src/mongo/dbtests/query_stage_count.cpp
@@ -292,11 +292,13 @@ public:
         if (interjection == 0) {
             // At this point, our first interjection, we've counted _locs[0]
             // and are about to count _locs[1]
+            WriteUnitOfWork wunit(&_txn);
             count_stage.invalidate(&_txn, _locs[interjection], INVALIDATION_DELETION);
             remove(_locs[interjection]);
 
             count_stage.invalidate(&_txn, _locs[interjection + 1], INVALIDATION_DELETION);
             remove(_locs[interjection + 1]);
+            wunit.commit();
         }
     }
 };
diff --git a/src/mongo/dbtests/query_stage_delete.cpp b/src/mongo/dbtests/query_stage_delete.cpp
index 4191c86b3d7..9607cef4612 100644
--- a/src/mongo/dbtests/query_stage_delete.cpp
+++ b/src/mongo/dbtests/query_stage_delete.cpp
@@ -162,7 +162,11 @@ public:
 
         // Remove locs[targetDocIndex];
         deleteStage.saveState();
-        deleteStage.invalidate(&_txn, locs[targetDocIndex], INVALIDATION_DELETION);
+        {
+            WriteUnitOfWork wunit(&_txn);
+            deleteStage.invalidate(&_txn, locs[targetDocIndex], INVALIDATION_DELETION);
+            wunit.commit();
+        }
         BSONObj targetDoc = coll->docFor(&_txn, locs[targetDocIndex]).value();
         ASSERT(!targetDoc.isEmpty());
         remove(targetDoc);
diff --git a/src/mongo/dbtests/query_stage_update.cpp b/src/mongo/dbtests/query_stage_update.cpp
index 96b6806fa49..4306af97848 100644
--- a/src/mongo/dbtests/query_stage_update.cpp
+++ b/src/mongo/dbtests/query_stage_update.cpp
@@ -304,7 +304,11 @@ public:
 
             // Remove locs[targetDocIndex];
             updateStage->saveState();
-            updateStage->invalidate(&_txn, locs[targetDocIndex], INVALIDATION_DELETION);
+            {
+                WriteUnitOfWork wunit(&_txn);
+                updateStage->invalidate(&_txn, locs[targetDocIndex], INVALIDATION_DELETION);
+                wunit.commit();
+            }
             BSONObj targetDoc = coll->docFor(&_txn, locs[targetDocIndex]).value();
             ASSERT(!targetDoc.isEmpty());
             remove(targetDoc);