SERVER-15187 Update saved OperationContext in all PlanStages

Every stage that has a field of type OperationContext*
should be updating it when ::restoreState(OperationContext*) is called.
Otherwise it is retaining a reference to deleted memory.

This bug hasn't surfaced before because:
1. the OperationContext used to be stack allocated, so it would have the
   same address every time
2. mmapv1 doesn't always need to dereference the OperationContext

Closes #781

Signed-off-by: Benety Goh <benety@mongodb.com>

19 files changed, 19 insertions, 0 deletions

diff --git a/src/mongo/db/exec/and_hash.cpp b/src/mongo/db/exec/and_hash.cpp
index fde17410232..9a196ae5ecc 100644
--- a/src/mongo/db/exec/and_hash.cpp
+++ b/src/mongo/db/exec/and_hash.cpp
@@ -445,6 +445,7 @@ namespace mongo {
     }
 
     void AndHashStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
 
         for (size_t i = 0; i < _children.size(); ++i) {
diff --git a/src/mongo/db/exec/and_sorted.cpp b/src/mongo/db/exec/and_sorted.cpp
index e81401a7bd3..57d7c1c2e83 100644
--- a/src/mongo/db/exec/and_sorted.cpp
+++ b/src/mongo/db/exec/and_sorted.cpp
@@ -268,6 +268,7 @@ namespace mongo {
     }
 
     void AndSortedStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
 
         for (size_t i = 0; i < _children.size(); ++i) {
diff --git a/src/mongo/db/exec/collection_scan.cpp b/src/mongo/db/exec/collection_scan.cpp
index 3ad6f519f0c..f05c3431790 100644
--- a/src/mongo/db/exec/collection_scan.cpp
+++ b/src/mongo/db/exec/collection_scan.cpp
@@ -156,6 +156,7 @@ namespace mongo {
     }
 
     void CollectionScan::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         if (NULL != _iter) {
             if (!_iter->restoreState(opCtx)) {
diff --git a/src/mongo/db/exec/count.cpp b/src/mongo/db/exec/count.cpp
index 077ec43db73..9f76ae10d1b 100644
--- a/src/mongo/db/exec/count.cpp
+++ b/src/mongo/db/exec/count.cpp
@@ -165,6 +165,7 @@ namespace mongo {
     }
 
     void CountStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         if (_child.get()) {
             _child->restoreState(opCtx);
diff --git a/src/mongo/db/exec/count_scan.cpp b/src/mongo/db/exec/count_scan.cpp
index daaed0fd545..4658ab9f759 100644
--- a/src/mongo/db/exec/count_scan.cpp
+++ b/src/mongo/db/exec/count_scan.cpp
@@ -155,6 +155,7 @@ namespace mongo {
     }
 
     void CountScan::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         if (_hitEnd || (NULL == _btreeCursor.get())) { return; }
 
diff --git a/src/mongo/db/exec/delete.cpp b/src/mongo/db/exec/delete.cpp
index a83d1b11dee..d7e689e62f6 100644
--- a/src/mongo/db/exec/delete.cpp
+++ b/src/mongo/db/exec/delete.cpp
@@ -153,6 +153,7 @@ namespace mongo {
     }
 
     void DeleteStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
     }
diff --git a/src/mongo/db/exec/distinct_scan.cpp b/src/mongo/db/exec/distinct_scan.cpp
index 933b65e6843..073a2ada59e 100644
--- a/src/mongo/db/exec/distinct_scan.cpp
+++ b/src/mongo/db/exec/distinct_scan.cpp
@@ -159,6 +159,7 @@ namespace mongo {
     }
 
     void DistinctScan::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
 
         if (_hitEnd || (NULL == _btreeCursor.get())) { return; }
diff --git a/src/mongo/db/exec/fetch.cpp b/src/mongo/db/exec/fetch.cpp
index 5e42d09e21d..62be74eda40 100644
--- a/src/mongo/db/exec/fetch.cpp
+++ b/src/mongo/db/exec/fetch.cpp
@@ -119,6 +119,7 @@ namespace mongo {
     }
 
     void FetchStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
     }
diff --git a/src/mongo/db/exec/group.cpp b/src/mongo/db/exec/group.cpp
index c53723cc772..117651ea3e3 100644
--- a/src/mongo/db/exec/group.cpp
+++ b/src/mongo/db/exec/group.cpp
@@ -230,6 +230,7 @@ namespace mongo {
     }
 
     void GroupStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
         return;
diff --git a/src/mongo/db/exec/idhack.cpp b/src/mongo/db/exec/idhack.cpp
index c82c7bb4de7..7b1845da8d2 100644
--- a/src/mongo/db/exec/idhack.cpp
+++ b/src/mongo/db/exec/idhack.cpp
@@ -135,6 +135,7 @@ namespace mongo {
     }
 
     void IDHackStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
     }
 
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index 1f78d7a7003..ffde08be60e 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -250,6 +250,7 @@ namespace mongo {
     }
 
     void IndexScan::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
 
         if (_hitEnd || (NULL == _indexCursor.get())) { return; }
diff --git a/src/mongo/db/exec/merge_sort.cpp b/src/mongo/db/exec/merge_sort.cpp
index 4b11e711843..f8b30416dce 100644
--- a/src/mongo/db/exec/merge_sort.cpp
+++ b/src/mongo/db/exec/merge_sort.cpp
@@ -192,6 +192,7 @@ namespace mongo {
     }
 
     void MergeSortStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         for (size_t i = 0; i < _children.size(); ++i) {
             _children[i]->restoreState(opCtx);
diff --git a/src/mongo/db/exec/multi_iterator.cpp b/src/mongo/db/exec/multi_iterator.cpp
index cda50f6e10a..8887687ddd4 100644
--- a/src/mongo/db/exec/multi_iterator.cpp
+++ b/src/mongo/db/exec/multi_iterator.cpp
@@ -68,6 +68,7 @@ namespace mongo {
     }
 
     void MultiIteratorStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         for (size_t i = 0; i < _iterators.size(); i++) {
             if (!_iterators[i]->restoreState(opCtx)) {
                 kill();
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
index 19c16bbc7d3..3836f10cf19 100644
--- a/src/mongo/db/exec/multi_plan.cpp
+++ b/src/mongo/db/exec/multi_plan.cpp
@@ -412,6 +412,7 @@ namespace mongo {
     }
 
     void MultiPlanStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         if (_failure) return;
 
         // this logic is from multi_plan_runner
diff --git a/src/mongo/db/exec/near.cpp b/src/mongo/db/exec/near.cpp
index 21aee43a7f7..53304381eb1 100644
--- a/src/mongo/db/exec/near.cpp
+++ b/src/mongo/db/exec/near.cpp
@@ -273,6 +273,7 @@ namespace mongo {
     }
 
     void NearStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_stats->common.unyields;
         if (_nextInterval) {
             _nextInterval->covering->restoreState(opCtx);
diff --git a/src/mongo/db/exec/oplogstart.cpp b/src/mongo/db/exec/oplogstart.cpp
index e2717e78512..4907b0e3aee 100644
--- a/src/mongo/db/exec/oplogstart.cpp
+++ b/src/mongo/db/exec/oplogstart.cpp
@@ -165,6 +165,7 @@ namespace mongo {
     }
 
     void OplogStart::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         if (_cs) {
             _cs->restoreState(opCtx);
         }
diff --git a/src/mongo/db/exec/sort.cpp b/src/mongo/db/exec/sort.cpp
index 22d60076104..02665c9b5b0 100644
--- a/src/mongo/db/exec/sort.cpp
+++ b/src/mongo/db/exec/sort.cpp
@@ -425,6 +425,7 @@ namespace mongo {
     }
 
     void SortStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
     }
diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp
index 8b9fb54ccb7..973c1d8a54b 100644
--- a/src/mongo/db/exec/subplan.cpp
+++ b/src/mongo/db/exec/subplan.cpp
@@ -438,6 +438,7 @@ namespace mongo {
     }
 
     void SubplanStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
         if (_killed) {
             return;
diff --git a/src/mongo/db/exec/text.cpp b/src/mongo/db/exec/text.cpp
index f040a274c01..91a696ab186 100644
--- a/src/mongo/db/exec/text.cpp
+++ b/src/mongo/db/exec/text.cpp
@@ -113,6 +113,7 @@ namespace mongo {
     }
 
     void TextStage::restoreState(OperationContext* opCtx) {
+        _txn = opCtx;
         ++_commonStats.unyields;
 
         for (size_t i = 0; i < _scanners.size(); ++i) {