SERVER-15541 add PlanExecutor yield policies and registrations

4 files changed, 97 insertions, 54 deletions

diff --git a/src/mongo/db/query/internal_plans.h b/src/mongo/db/query/internal_plans.h
index b623e06aee5..c16aa251f2d 100644
--- a/src/mongo/db/query/internal_plans.h
+++ b/src/mongo/db/query/internal_plans.h
@@ -90,10 +90,8 @@ namespace mongo {
             }
 
             CollectionScan* cs = new CollectionScan(txn, params, ws, NULL);
-            PlanExecutor* exec = new PlanExecutor(txn, ws, cs, collection);
-            // 'exec' will be registered until it is destroyed.
-            exec->registerExecInternalPlan();
-            return exec;
+            // Takes ownership of 'ws' and 'cs'.
+            return new PlanExecutor(txn, ws, cs, collection);
         }
 
         /**
@@ -125,10 +123,7 @@ namespace mongo {
                 root = new FetchStage(txn, ws, root, NULL, collection);
             }
 
-            PlanExecutor* exec = new PlanExecutor(txn, ws, root, collection);
-            // 'exec' will be registered until it is destroyed.
-            exec->registerExecInternalPlan();
-            return exec;
+            return new PlanExecutor(txn, ws, root, collection);
         }
     };
 
diff --git a/src/mongo/db/query/new_find.cpp b/src/mongo/db/query/new_find.cpp
index 6c8682c8851..3898424d031 100644
--- a/src/mongo/db/query/new_find.cpp
+++ b/src/mongo/db/query/new_find.cpp
@@ -442,8 +442,7 @@ namespace mongo {
         OplogStart* stage = new OplogStart(txn, collection, tsExpr, oplogws);
 
         // Takes ownership of ws and stage.
-        auto_ptr<PlanExecutor> exec(new PlanExecutor(txn, oplogws, stage, collection));
-        exec->registerExecInternalPlan();
+        scoped_ptr<PlanExecutor> exec(new PlanExecutor(txn, oplogws, stage, collection));
 
         // The stage returns a DiskLoc of where to start.
         DiskLoc startLoc;
@@ -596,6 +595,11 @@ namespace mongo {
         verify(NULL != rawExec);
         auto_ptr<PlanExecutor> exec(rawExec);
 
+        // We want the PlanExecutor to yield automatically, but we handle registration of the
+        // executor ourselves. We want to temporarily register the executor while we are generating
+        // this batch of results, and then unregister and re-register with ClientCursor for getmore.
+        exec->setYieldPolicy(PlanExecutor::YIELD_AUTO);
+
         // If it's actually an explain, do the explain and return rather than falling through
         // to the normal query execution loop.
         if (pq.isExplain()) {
@@ -676,9 +680,6 @@ namespace mongo {
         // Do we save the PlanExecutor in a ClientCursor for getMore calls later?
         bool saveClientCursor = false;
 
-        // The executor registers itself with the active executors list in ClientCursor.
-        auto_ptr<ScopedExecutorRegistration> safety(new ScopedExecutorRegistration(exec.get()));
-
         BSONObj obj;
         PlanExecutor::ExecState state;
         // uint64_t numMisplacedDocs = 0;
@@ -728,7 +729,7 @@ namespace mongo {
         // If we don't cache the executor later, we are deleting it, so it must be deregistered.
         //
         // So, no matter what, deregister the executor.
-        safety.reset();
+        exec->deregisterExec();
 
         // Caller expects exceptions thrown in certain cases.
         if (PlanExecutor::EXEC_ERROR == state) {
diff --git a/src/mongo/db/query/plan_executor.cpp b/src/mongo/db/query/plan_executor.cpp
index af1b5f2f4aa..601091a567b 100644
--- a/src/mongo/db/query/plan_executor.cpp
+++ b/src/mongo/db/query/plan_executor.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/exec/plan_stats.h"
 #include "mongo/db/exec/working_set.h"
 #include "mongo/db/exec/working_set_common.h"
+#include "mongo/db/query/plan_yield_policy.h"
 
 #include "mongo/util/stacktrace.h"
 
@@ -251,10 +252,14 @@ namespace mongo {
         return _killed || _root->isEOF();
     }
 
-    void PlanExecutor::registerExecInternalPlan() {
+    void PlanExecutor::registerExec() {
         _safety.reset(new ScopedExecutorRegistration(this));
     }
 
+    void PlanExecutor::deregisterExec() {
+        _safety.reset();
+    }
+
     void PlanExecutor::kill() {
         _killed = true;
         _collection = NULL;
@@ -298,24 +303,40 @@ namespace mongo {
     }
 
     void PlanExecutor::setYieldPolicy(YieldPolicy policy, bool registerExecutor) {
-        // TODO: implement.
+        if (PlanExecutor::YIELD_MANUAL == policy) {
+            _yieldPolicy.reset();
+        }
+        else {
+            invariant(PlanExecutor::YIELD_AUTO == policy);
+            _yieldPolicy.reset(new PlanYieldPolicy());
+
+            // Runners that yield automatically generally need to be registered so that
+            // after yielding, they receive notifications of events like deletions and
+            // index drops. The only exception is that a few PlanExecutors get registered
+            // by ClientCursor instead of being registered here.
+            if (registerExecutor) {
+                this->registerExec();
+            }
+        }
     }
 
     //
     // ScopedExecutorRegistration
     //
 
-    ScopedExecutorRegistration::ScopedExecutorRegistration(PlanExecutor* exec)
+    PlanExecutor::ScopedExecutorRegistration::ScopedExecutorRegistration(PlanExecutor* exec)
         : _exec(exec) {
         // Collection can be null for an EOFStage plan, or other places where registration
         // is not needed.
-        if ( _exec->collection() )
-            _exec->collection()->cursorCache()->registerExecutor( exec );
+        if (_exec->collection()) {
+            _exec->collection()->cursorCache()->registerExecutor(exec);
+        }
     }
 
-    ScopedExecutorRegistration::~ScopedExecutorRegistration() {
-        if ( _exec->collection() )
-            _exec->collection()->cursorCache()->deregisterExecutor( _exec );
+    PlanExecutor::ScopedExecutorRegistration::~ScopedExecutorRegistration() {
+        if (_exec->collection()) {
+            _exec->collection()->cursorCache()->deregisterExecutor(_exec);
+        }
     }
 
 } // namespace mongo
diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h
index a52136abab5..db616fc7281 100644
--- a/src/mongo/db/query/plan_executor.h
+++ b/src/mongo/db/query/plan_executor.h
@@ -42,27 +42,10 @@ namespace mongo {
     class PlanStage;
     class PlanExecutor;
     struct PlanStageStats;
+    class PlanYieldPolicy;
     class WorkingSet;
 
     /**
-     * RAII approach to ensuring that plan executors are deregistered.
-     *
-     * While retrieving the first batch of results, newRunQuery manually registers the executor with
-     * ClientCursor.  Certain query execution paths, namely $where, can throw an exception.  If we
-     * fail to deregister the executor, we will call invalidate/kill on the
-     * still-registered-yet-deleted executor.
-     *
-     * For any subsequent calls to getMore, the executor is already registered with ClientCursor
-     * by virtue of being cached, so this exception-proofing is not required.
-     */
-    struct ScopedExecutorRegistration {
-        ScopedExecutorRegistration(PlanExecutor* exec);
-        ~ScopedExecutorRegistration();
-
-        PlanExecutor* const _exec;
-    };
-
-    /**
      * A PlanExecutor is the abstraction that knows how to crank a tree of stages into execution.
      * The executor is usually part of a larger abstraction that is interacting with the cache
      * and/or the query optimizer.
@@ -101,7 +84,27 @@ namespace mongo {
             YIELD_AUTO,
 
             // Owner must yield manually if yields are requested.  How to yield yourself:
-            // XXX: go on about this
+            //
+            // 0. Let's say you have PlanExecutor* exec.
+            //
+            // 1. Register your PlanExecutor with ClientCursor. Registered executors are informed
+            // about DiskLoc deletions and namespace invalidation, as well as other important
+            // events. Do this either by  calling registerExec() on the executor. This can be done
+            // once you get your executor, or could be done per-yield.
+            //
+            // 2. Call exec->saveState() before you yield.
+            //
+            // 3. Call Yield::yieldAllLocks(), passing in the executor's OperationContext*. This
+            // causes the executor to give up its locks and block so that it goes to the back of
+            // the scheduler's queue.
+            //
+            // 4. Call exec->restoreState() before using the executor again.
+            //
+            // 5. The next call to exec->getNext() may return DEAD.
+            //
+            // 6. Make sure the executor gets deregistered from ClientCursor. PlanExecutor does
+            // this in an RAII fashion when it is destroyed, or you can explicity call
+            // unregisterExec() on the PlanExecutor.
             YIELD_MANUAL,
         };
 
@@ -251,14 +254,14 @@ namespace mongo {
          * receives notifications for events that happen while yielding any locks.
          *
          * Deregistration happens automatically when this plan executor is destroyed.
-         *
-         * Used for internal clients of the query system:
-         *  -- InternalPlanner::collectionScan(...) (see internal_plans.h)
-         *  -- InternalPlanner::indexScan(...) (see internal_plans.h)
-         *  -- getOplogStartHack(...) (see new_find.cpp)
-         *  -- storeCurrentLocs(...) (see d_migrate.cpp)
          */
-        void registerExecInternalPlan();
+        void registerExec();
+
+        /**
+         * Unregister this PlanExecutor. Normally you want the PlanExecutor to be registered
+         * for its lifetime, and you shouldn't have to call this explicitly.
+         */
+        void deregisterExec();
 
         /**
          * If we're yielding locks, the database we're operating over or any collection we're
@@ -281,18 +284,37 @@ namespace mongo {
         static std::string statestr(ExecState s);
 
         /**
-         * Change the yield policy of the PlanExecutor to 'policy'.
+         * Change the yield policy of the PlanExecutor to 'policy'. If 'registerExecutor' is true,
+         * and the yield policy is YIELD_AUTO, then the plan executor gets registered to receive
+         * notifications of events from other threads.
          *
-         * XXX: everybody who sets the policy to YIELD_AUTO really wants to call
-         * registerExecInternalPlan immediately after EXCEPT new_find.cpp...so we expose
-         * the ability to register (or not) via registerExecutor, rather than require
-         * all users to have yet another RAII object.  Only new_find.cpp should ever
-         * set registerExecutor to false.
+         * Everybody who sets the policy to YIELD_AUTO really wants to call registerExec()
+         * immediately after EXCEPT commands that create cursors...so we expose the ability to
+         * register (or not) here, rather than require all users to have yet another RAII object.
+         * Only cursor-creating things like new_find.cpp set registerExecutor to false.
          */
         void setYieldPolicy(YieldPolicy policy, bool registerExecutor = true);
 
     private:
         /**
+         * RAII approach to ensuring that plan executors are deregistered.
+         *
+         * While retrieving the first batch of results, newRunQuery manually registers the executor
+         * with ClientCursor.  Certain query execution paths, namely $where, can throw an exception.
+         * If we fail to deregister the executor, we will call invalidate/kill on the
+         * still-registered-yet-deleted executor.
+         *
+         * For any subsequent calls to getMore, the executor is already registered with ClientCursor
+         * by virtue of being cached, so this exception-proofing is not required.
+         */
+        struct ScopedExecutorRegistration {
+            ScopedExecutorRegistration(PlanExecutor* exec);
+            ~ScopedExecutorRegistration();
+
+            PlanExecutor* const _exec;
+        };
+
+        /**
          * Initialize the namespace using either the canonical query or the collection.
          */
         void initNs();
@@ -319,6 +341,10 @@ namespace mongo {
         // Did somebody drop an index we care about or the namespace we're looking at?  If so,
         // we'll be killed.
         bool _killed;
+
+        // If the yield policy is YIELD_AUTO, this is used to enforce automatic yielding. The plan
+        // may yield on any call to getNext() if this is non-NULL.
+        boost::scoped_ptr<PlanYieldPolicy> _yieldPolicy;
     };
 
 }  // namespace mongo