SERVER-19364 move query stage OperationContext pointer management into the base class

Author: coollog

src/mongo/db/commands/list_collections.cpp

@@ -45,13 +45,15 @@
 #include "mongo/db/query/find_constants.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/storage/storage_engine.h"
+#include "mongo/stdx/memory.h"
 
 namespace mongo {
 
-using std::unique_ptr;
 using std::list;
 using std::string;
 using std::stringstream;
+using std::unique_ptr;
+using stdx::make_unique;
 
 class CmdListCollections : public Command {
 public:
@@ -99,7 +101,7 @@ class CmdListCollections : public Command {
              int,
              string& errmsg,
              BSONObjBuilder& result) {
-        std::unique_ptr<MatchExpression> matcher;
+        unique_ptr<MatchExpression> matcher;
         BSONElement filterElt = jsobj["filter"];
         if (!filterElt.eoo()) {
             if (filterElt.type() != mongo::Object) {
@@ -134,8 +136,8 @@ class CmdListCollections : public Command {
             names.sort();
         }
 
-        std::unique_ptr<WorkingSet> ws(new WorkingSet());
-        std::unique_ptr<QueuedDataStage> root(new QueuedDataStage(ws.get()));
+        auto ws = make_unique<WorkingSet>();
+        auto root = make_unique<QueuedDataStage>(txn, ws.get());
 
         for (std::list<std::string>::const_iterator i = names.begin(); i != names.end(); ++i) {
             const std::string& ns = *i;
@@ -175,7 +177,7 @@ class CmdListCollections : public Command {
         if (!statusWithPlanExecutor.isOK()) {
             return appendCommandStatus(result, statusWithPlanExecutor.getStatus());
         }
-        std::unique_ptr<PlanExecutor> exec = std::move(statusWithPlanExecutor.getValue());
+        unique_ptr<PlanExecutor> exec = std::move(statusWithPlanExecutor.getValue());
 
         BSONArrayBuilder firstBatch;
 

src/mongo/db/commands/list_indexes.cpp

@@ -45,12 +45,15 @@
 #include "mongo/db/query/find_constants.h"
 #include "mongo/db/service_context.h"
 #include "mongo/db/storage/storage_engine.h"
+#include "mongo/stdx/memory.h"
 
 namespace mongo {
 
 using std::string;
 using std::stringstream;
+using std::unique_ptr;
 using std::vector;
+using stdx::make_unique;
 
 /**
  * Lists the indexes for a given collection.
@@ -143,8 +146,8 @@ class CmdListIndexes : public Command {
         }
         MONGO_WRITE_CONFLICT_RETRY_LOOP_END(txn, "listIndexes", ns.ns());
 
-        std::unique_ptr<WorkingSet> ws(new WorkingSet());
-        std::unique_ptr<QueuedDataStage> root(new QueuedDataStage(ws.get()));
+        auto ws = make_unique<WorkingSet>();
+        auto root = make_unique<QueuedDataStage>(txn, ws.get());
 
         for (size_t i = 0; i < indexNames.size(); i++) {
             BSONObj indexSpec;
@@ -173,7 +176,7 @@ class CmdListIndexes : public Command {
         if (!statusWithPlanExecutor.isOK()) {
             return appendCommandStatus(result, statusWithPlanExecutor.getStatus());
         }
-        std::unique_ptr<PlanExecutor> exec = std::move(statusWithPlanExecutor.getValue());
+        unique_ptr<PlanExecutor> exec = std::move(statusWithPlanExecutor.getValue());
 
         BSONArrayBuilder firstBatch;
 

src/mongo/db/commands/pipeline_command.cpp

@@ -51,6 +51,7 @@
 #include "mongo/db/query/find_constants.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/storage_options.h"
+#include "mongo/stdx/memory.h"
 
 namespace mongo {
 
@@ -60,6 +61,7 @@ using std::shared_ptr;
 using std::string;
 using std::stringstream;
 using std::unique_ptr;
+using stdx::make_unique;
 
 /**
  * Returns true if we need to keep a ClientCursor saved for this pipeline (for future getMore
@@ -230,9 +232,8 @@ class PipelineCommand : public Command {
             // Create the PlanExecutor which returns results from the pipeline. The WorkingSet
             // ('ws') and the PipelineProxyStage ('proxy') will be owned by the created
             // PlanExecutor.
-            unique_ptr<WorkingSet> ws(new WorkingSet());
-            unique_ptr<PipelineProxyStage> proxy(
-                new PipelineProxyStage(pPipeline, input, ws.get()));
+            auto ws = make_unique<WorkingSet>();
+            auto proxy = make_unique<PipelineProxyStage>(txn, pPipeline, input, ws.get());
 
             auto statusWithPlanExecutor = (NULL == collection)
                 ? PlanExecutor::make(

src/mongo/db/exec/and_hash.cpp

@@ -54,17 +54,20 @@ const size_t AndHashStage::kLookAheadWorks = 10;
 // static
 const char* AndHashStage::kStageType = "AND_HASH";
 
-AndHashStage::AndHashStage(WorkingSet* ws, const Collection* collection)
-    : PlanStage(kStageType),
+AndHashStage::AndHashStage(OperationContext* opCtx, WorkingSet* ws, const Collection* collection)
+    : PlanStage(kStageType, opCtx),
       _collection(collection),
       _ws(ws),
       _hashingChildren(true),
       _currentChild(0),
       _memUsage(0),
       _maxMemUsage(kDefaultMaxMemUsageBytes) {}
 
-AndHashStage::AndHashStage(WorkingSet* ws, const Collection* collection, size_t maxMemUsage)
-    : PlanStage(kStageType),
+AndHashStage::AndHashStage(OperationContext* opCtx,
+                           WorkingSet* ws,
+                           const Collection* collection,
+                           size_t maxMemUsage)
+    : PlanStage(kStageType, opCtx),
       _collection(collection),
       _ws(ws),
       _hashingChildren(true),

src/mongo/db/exec/and_hash.h

@@ -51,12 +51,15 @@ namespace mongo {
  */
 class AndHashStage final : public PlanStage {
 public:
-    AndHashStage(WorkingSet* ws, const Collection* collection);
+    AndHashStage(OperationContext* opCtx, WorkingSet* ws, const Collection* collection);
 
     /**
      * For testing only. Allows tests to set memory usage threshold.
      */
-    AndHashStage(WorkingSet* ws, const Collection* collection, size_t maxMemUsage);
+    AndHashStage(OperationContext* opCtx,
+                 WorkingSet* ws,
+                 const Collection* collection,
+                 size_t maxMemUsage);
 
     void addChild(PlanStage* child);
 

src/mongo/db/exec/and_sorted.cpp

@@ -44,8 +44,10 @@ using stdx::make_unique;
 // static
 const char* AndSortedStage::kStageType = "AND_SORTED";
 
-AndSortedStage::AndSortedStage(WorkingSet* ws, const Collection* collection)
-    : PlanStage(kStageType),
+AndSortedStage::AndSortedStage(OperationContext* opCtx,
+                               WorkingSet* ws,
+                               const Collection* collection)
+    : PlanStage(kStageType, opCtx),
       _collection(collection),
       _ws(ws),
       _targetNode(numeric_limits<size_t>::max()),

src/mongo/db/exec/and_sorted.h

@@ -53,7 +53,7 @@ namespace mongo {
  */
 class AndSortedStage final : public PlanStage {
 public:
-    AndSortedStage(WorkingSet* ws, const Collection* collection);
+    AndSortedStage(OperationContext* opCtx, WorkingSet* ws, const Collection* collection);
 
     void addChild(PlanStage* child);
 

src/mongo/db/exec/cached_plan.cpp

@@ -60,8 +60,7 @@ CachedPlanStage::CachedPlanStage(OperationContext* txn,
                                  const QueryPlannerParams& params,
                                  size_t decisionWorks,
                                  PlanStage* root)
-    : PlanStage(kStageType),
-      _txn(txn),
+    : PlanStage(kStageType, txn),
       _collection(collection),
       _ws(ws),
       _canonicalQuery(cq),
@@ -228,15 +227,16 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache) {
 
         PlanStage* newRoot;
         // Only one possible plan. Build the stages from the solution.
-        verify(StageBuilder::build(_txn, _collection, *solutions[0], _ws, &newRoot));
+        verify(StageBuilder::build(getOpCtx(), _collection, *solutions[0], _ws, &newRoot));
         _children.emplace_back(newRoot);
         _replannedQs.reset(solutions.popAndReleaseBack());
         return Status::OK();
     }
 
     // Many solutions. Create a MultiPlanStage to pick the best, update the cache,
     // and so on. The working set will be shared by all candidate plans.
-    _children.emplace_back(new MultiPlanStage(_txn, _collection, _canonicalQuery, shouldCache));
+    _children.emplace_back(
+        new MultiPlanStage(getOpCtx(), _collection, _canonicalQuery, shouldCache));
     MultiPlanStage* multiPlanStage = static_cast<MultiPlanStage*>(child().get());
 
     for (size_t ix = 0; ix < solutions.size(); ++ix) {
@@ -245,7 +245,7 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache) {
         }
 
         PlanStage* nextPlanRoot;
-        verify(StageBuilder::build(_txn, _collection, *solutions[ix], _ws, &nextPlanRoot));
+        verify(StageBuilder::build(getOpCtx(), _collection, *solutions[ix], _ws, &nextPlanRoot));
 
         // Takes ownership of 'solutions[ix]' and 'nextPlanRoot'.
         multiPlanStage->addPlan(solutions.releaseAt(ix), nextPlanRoot, _ws);
@@ -291,11 +291,6 @@ PlanStage::StageState CachedPlanStage::work(WorkingSetID* out) {
     return childStatus;
 }
 
-
-void CachedPlanStage::doReattachToOperationContext(OperationContext* opCtx) {
-    _txn = opCtx;
-}
-
 void CachedPlanStage::doInvalidate(OperationContext* txn,
                                    const RecordId& dl,
                                    InvalidationType type) {

src/mongo/db/exec/cached_plan.h

@@ -65,7 +65,6 @@ class CachedPlanStage final : public PlanStage {
 
     StageState work(WorkingSetID* out) final;
 
-    void doReattachToOperationContext(OperationContext* opCtx) final;
     void doInvalidate(OperationContext* txn, const RecordId& dl, InvalidationType type) final;
 
     StageType stageType() const final {
@@ -115,9 +114,6 @@ class CachedPlanStage final : public PlanStage {
      */
     Status tryYield(PlanYieldPolicy* yieldPolicy);
 
-    // Not owned.
-    OperationContext* _txn;
-
     // Not owned. Must be non-null.
     Collection* _collection;
 

src/mongo/db/exec/collection_scan.cpp

@@ -58,8 +58,7 @@ CollectionScan::CollectionScan(OperationContext* txn,
                                const CollectionScanParams& params,
                                WorkingSet* workingSet,
                                const MatchExpression* filter)
-    : PlanStage(kStageType),
-      _txn(txn),
+    : PlanStage(kStageType, txn),
       _workingSet(workingSet),
       _filter(filter),
       _params(params),
@@ -98,7 +97,7 @@ PlanStage::StageState CollectionScan::work(WorkingSetID* out) {
     try {
         if (needToMakeCursor) {
             const bool forward = _params.direction == CollectionScanParams::FORWARD;
-            _cursor = _params.collection->getCursor(_txn, forward);
+            _cursor = _params.collection->getCursor(getOpCtx(), forward);
 
             if (!_lastSeenId.isNull()) {
                 invariant(_params.tailable);
@@ -165,7 +164,7 @@ PlanStage::StageState CollectionScan::work(WorkingSetID* out) {
     WorkingSetID id = _workingSet->allocate();
     WorkingSetMember* member = _workingSet->get(id);
     member->loc = record->id;
-    member->obj = {_txn->recoveryUnit()->getSnapshotId(), record->data.releaseToBson()};
+    member->obj = {getOpCtx()->recoveryUnit()->getSnapshotId(), record->data.releaseToBson()};
     _workingSet->transitionToLocAndObj(id);
 
     return returnIfMatches(member, id, out);
@@ -223,23 +222,21 @@ void CollectionScan::doSaveState() {
 void CollectionScan::doRestoreState() {
     if (_cursor) {
         if (!_cursor->restore()) {
-            warning() << "Could not restore RecordCursor for CollectionScan: " << _txn->getNS();
+            warning() << "Could not restore RecordCursor for CollectionScan: "
+                      << getOpCtx()->getNS();
             _isDead = true;
         }
     }
 }
 
 void CollectionScan::doDetachFromOperationContext() {
-    _txn = NULL;
     if (_cursor)
         _cursor->detachFromOperationContext();
 }
 
-void CollectionScan::doReattachToOperationContext(OperationContext* opCtx) {
-    invariant(_txn == NULL);
-    _txn = opCtx;
+void CollectionScan::doReattachToOperationContext() {
     if (_cursor)
-        _cursor->reattachToOperationContext(opCtx);
+        _cursor->reattachToOperationContext(getOpCtx());
 }
 
 unique_ptr<PlanStageStats> CollectionScan::getStats() {