SERVER-15675 Stages should clear their OperationContext in saveState()

While state is saved, threads pass in their own OperationContext to
invalidate().  When state is restored, the stage will resume with the
OperationContext passed in to restoreState().

18 files changed, 44 insertions, 4 deletions

diff --git a/src/mongo/db/exec/collection_scan.cpp b/src/mongo/db/exec/collection_scan.cpp
index b0c1267c8c4..7539ea96289 100644
--- a/src/mongo/db/exec/collection_scan.cpp
+++ b/src/mongo/db/exec/collection_scan.cpp
@@ -205,6 +205,7 @@ namespace mongo {
     }
 
     void CollectionScan::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         if (NULL != _iter) {
             _iter->saveState();
@@ -212,6 +213,7 @@ namespace mongo {
     }
 
     void CollectionScan::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         if (NULL != _iter) {
diff --git a/src/mongo/db/exec/count.cpp b/src/mongo/db/exec/count.cpp
index 6bcd3d63ce6..58cf1bec598 100644
--- a/src/mongo/db/exec/count.cpp
+++ b/src/mongo/db/exec/count.cpp
@@ -164,6 +164,7 @@ namespace mongo {
     }
 
     void CountStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         if (_child.get()) {
             _child->saveState();
@@ -171,6 +172,7 @@ namespace mongo {
     }
 
     void CountStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         if (_child.get()) {
diff --git a/src/mongo/db/exec/count_scan.cpp b/src/mongo/db/exec/count_scan.cpp
index cff8b4ccaaa..09d4b098d1d 100644
--- a/src/mongo/db/exec/count_scan.cpp
+++ b/src/mongo/db/exec/count_scan.cpp
@@ -148,6 +148,7 @@ namespace mongo {
     }
 
     void CountScan::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         if (_hitEnd || (NULL == _btreeCursor.get())) { return; }
 
@@ -156,6 +157,7 @@ namespace mongo {
     }
 
     void CountScan::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _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 d13b8bf2ff8..d1244f1219b 100644
--- a/src/mongo/db/exec/delete.cpp
+++ b/src/mongo/db/exec/delete.cpp
@@ -96,7 +96,7 @@ namespace mongo {
 
             // TODO: Do we want to buffer docs and delete them in a group rather than
             // saving/restoring state repeatedly?
-            saveState();
+            _child->saveState();
 
             {
                 WriteUnitOfWork wunit(_txn);
@@ -128,7 +128,7 @@ namespace mongo {
             //  As restoreState may restore (recreate) cursors, cursors are tied to the
             //  transaction in which they are created, and a WriteUnitOfWork is a
             //  transaction, make sure to restore the state outside of the WritUnitOfWork.
-            restoreState(_txn);
+            _child->restoreState(_txn);
 
             ++_specificStats.docsDeleted;
 
@@ -159,11 +159,13 @@ namespace mongo {
     }
 
     void DeleteStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         _child->saveState();
     }
 
     void DeleteStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _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 ddbe6a49103..5f3f516cea0 100644
--- a/src/mongo/db/exec/distinct_scan.cpp
+++ b/src/mongo/db/exec/distinct_scan.cpp
@@ -170,6 +170,7 @@ namespace mongo {
     }
 
     void DistinctScan::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
 
         if (HIT_END == _scanState || INITIALIZING == _scanState) { return; }
@@ -183,6 +184,7 @@ namespace mongo {
     }
 
     void DistinctScan::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
 
diff --git a/src/mongo/db/exec/fetch.cpp b/src/mongo/db/exec/fetch.cpp
index 112bcbec1d3..bc1fd838385 100644
--- a/src/mongo/db/exec/fetch.cpp
+++ b/src/mongo/db/exec/fetch.cpp
@@ -152,11 +152,13 @@ namespace mongo {
     }
 
     void FetchStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         _child->saveState();
     }
 
     void FetchStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
diff --git a/src/mongo/db/exec/group.cpp b/src/mongo/db/exec/group.cpp
index af6b5958384..90a28c8f3f7 100644
--- a/src/mongo/db/exec/group.cpp
+++ b/src/mongo/db/exec/group.cpp
@@ -257,11 +257,13 @@ namespace mongo {
     }
 
     void GroupStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         _child->saveState();
     }
 
     void GroupStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         _child->restoreState(opCtx);
diff --git a/src/mongo/db/exec/idhack.cpp b/src/mongo/db/exec/idhack.cpp
index 261bfeac369..ffa727ec808 100644
--- a/src/mongo/db/exec/idhack.cpp
+++ b/src/mongo/db/exec/idhack.cpp
@@ -173,10 +173,12 @@ namespace mongo {
     }
 
     void IDHackStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
     }
 
     void IDHackStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
     }
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index b6cb55d54d9..06e821f14ed 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -242,6 +242,7 @@ namespace mongo {
     }
 
     void IndexScan::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
 
         if (HIT_END == _scanState || INITIALIZING == _scanState) { return; }
@@ -253,6 +254,7 @@ namespace mongo {
     }
 
     void IndexScan::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
 
diff --git a/src/mongo/db/exec/multi_iterator.cpp b/src/mongo/db/exec/multi_iterator.cpp
index fa3d10fbbb8..6dab3080da4 100644
--- a/src/mongo/db/exec/multi_iterator.cpp
+++ b/src/mongo/db/exec/multi_iterator.cpp
@@ -96,12 +96,14 @@ namespace mongo {
     }
 
     void MultiIteratorStage::saveState() {
+        _txn = NULL;
         for (size_t i = 0; i < _iterators.size(); i++) {
             _iterators[i]->saveState();
         }
     }
 
     void MultiIteratorStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         for (size_t i = 0; i < _iterators.size(); i++) {
             if (!_iterators[i]->restoreState(opCtx)) {
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
index 04488cfdec3..fc28e914acf 100644
--- a/src/mongo/db/exec/multi_plan.cpp
+++ b/src/mongo/db/exec/multi_plan.cpp
@@ -407,12 +407,14 @@ namespace mongo {
     }
 
     void MultiPlanStage::saveState() {
+        _txn = NULL;
         for (size_t i = 0; i < _candidates.size(); ++i) {
             _candidates[i].root->saveState();
         }
     }
 
     void MultiPlanStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
 
         for (size_t i = 0; i < _candidates.size(); ++i) {
diff --git a/src/mongo/db/exec/near.cpp b/src/mongo/db/exec/near.cpp
index 52e4c5858e4..294808970bd 100644
--- a/src/mongo/db/exec/near.cpp
+++ b/src/mongo/db/exec/near.cpp
@@ -303,6 +303,7 @@ namespace mongo {
     }
 
     void NearStage::saveState() {
+        _txn = NULL;
         ++_stats->common.yields;
         for (size_t i = 0; i < _childrenIntervals.size(); i++) {
             _childrenIntervals[i]->covering->saveState();
@@ -310,6 +311,7 @@ namespace mongo {
     }
 
     void NearStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_stats->common.unyields;
         for (size_t i = 0; i < _childrenIntervals.size(); i++) {
diff --git a/src/mongo/db/exec/oplogstart.cpp b/src/mongo/db/exec/oplogstart.cpp
index 5d33bd4f4ba..c73eaeb3087 100644
--- a/src/mongo/db/exec/oplogstart.cpp
+++ b/src/mongo/db/exec/oplogstart.cpp
@@ -156,6 +156,7 @@ namespace mongo {
     }
 
     void OplogStart::saveState() {
+        _txn = NULL;
         if (_cs) {
             _cs->saveState();
         }
@@ -166,6 +167,7 @@ namespace mongo {
     }
 
     void OplogStart::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         if (_cs) {
             _cs->restoreState(opCtx);
diff --git a/src/mongo/db/exec/pipeline_proxy.cpp b/src/mongo/db/exec/pipeline_proxy.cpp
index dce1651015f..69c18d5c060 100644
--- a/src/mongo/db/exec/pipeline_proxy.cpp
+++ b/src/mongo/db/exec/pipeline_proxy.cpp
@@ -95,6 +95,7 @@ namespace mongo {
     }
 
     void PipelineProxyStage::restoreState(OperationContext* opCtx) {
+        invariant(_pipeline->getContext()->opCtx == NULL);
         _pipeline->getContext()->opCtx = opCtx;
     }
 
diff --git a/src/mongo/db/exec/plan_stage.h b/src/mongo/db/exec/plan_stage.h
index 71432c06efa..1adba4bd906 100644
--- a/src/mongo/db/exec/plan_stage.h
+++ b/src/mongo/db/exec/plan_stage.h
@@ -201,6 +201,9 @@ namespace mongo {
          * any saved state and be ready to handle calls to work().
          *
          * Can only be called after saveState.
+         *
+         * If the stage needs an OperationContext during its execution, it may keep a handle to the
+         * provided OperationContext (which is valid until the next call to saveState()).
          */
         virtual void restoreState(OperationContext* opCtx) = 0;
 
@@ -210,6 +213,10 @@ namespace mongo {
          * DiskLoc.
          *
          * Can only be called after a saveState but before a restoreState.
+         *
+         * The provided OperationContext should be used if any work needs to be performed during the
+         * invalidate (as the state of the stage must be saved before any calls to invalidate, the
+         * stage's own OperationContext is inactive during the invalidate and should not be used).
          */
         virtual void invalidate(OperationContext* txn,
                                 const DiskLoc& dl,
diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp
index a8d27dd3faf..7c177377c8a 100644
--- a/src/mongo/db/exec/subplan.cpp
+++ b/src/mongo/db/exec/subplan.cpp
@@ -490,6 +490,7 @@ namespace mongo {
     }
 
     void SubplanStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         if (_killed) {
             return;
@@ -503,6 +504,7 @@ namespace mongo {
     }
 
     void SubplanStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         if (_killed) {
diff --git a/src/mongo/db/exec/text.cpp b/src/mongo/db/exec/text.cpp
index 60441eb2a81..d06f302e43b 100644
--- a/src/mongo/db/exec/text.cpp
+++ b/src/mongo/db/exec/text.cpp
@@ -109,6 +109,7 @@ namespace mongo {
     }
 
     void TextStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
 
         for (size_t i = 0; i < _scanners.size(); ++i) {
@@ -117,6 +118,7 @@ namespace mongo {
     }
 
     void TextStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
 
diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index 5cdc0e9aeba..dcba59dd5a6 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -741,7 +741,7 @@ namespace mongo {
             ++_specificStats.nMatched;
 
             // Save state before making changes
-            saveState();
+            _child->saveState();
 
             // Do the update and return.
             BSONObj reFetched;
@@ -780,7 +780,7 @@ namespace mongo {
             // As restoreState may restore (recreate) cursors, make sure to restore the
             // state outside of the WritUnitOfWork.
 
-            restoreState(_txn);
+            _child->restoreState(_txn);
 
             ++_commonStats.needTime;
             return PlanStage::NEED_TIME;
@@ -815,6 +815,7 @@ namespace mongo {
     }
 
     void UpdateStage::saveState() {
+        _txn = NULL;
         ++_commonStats.yields;
         _child->saveState();
     }
@@ -848,6 +849,7 @@ namespace mongo {
     }
 
     void UpdateStage::restoreState(OperationContext* opCtx) {
+        invariant(_txn == NULL);
         _txn = opCtx;
         ++_commonStats.unyields;
         // Restore our child.