SERVER-30410 Ensure executor is saved after tailable cursor time out.

6 files changed, 1408 insertions, 1149 deletions

diff --git a/src/mongo/dbtests/documentsourcetests.cpp b/src/mongo/dbtests/documentsourcetests.cpp
index 4af64c6be73..f5bd327705f 100644
--- a/src/mongo/dbtests/documentsourcetests.cpp
+++ b/src/mongo/dbtests/documentsourcetests.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/client.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbdirectclient.h"
+#include "mongo/db/exec/collection_scan.h"
 #include "mongo/db/exec/multi_plan.h"
 #include "mongo/db/exec/plan_stage.h"
 #include "mongo/db/pipeline/dependencies.h"
@@ -43,12 +44,15 @@
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/pipeline/pipeline.h"
 #include "mongo/db/query/get_executor.h"
+#include "mongo/db/query/mock_yield_policies.h"
 #include "mongo/db/query/plan_executor.h"
 #include "mongo/db/query/query_planner.h"
 #include "mongo/db/query/stage_builder.h"
 #include "mongo/dbtests/dbtests.h"
+#include "mongo/util/scopeguard.h"
 
-namespace DocumentSourceCursorTests {
+namespace mongo {
+namespace {
 
 using boost::intrusive_ptr;
 using std::unique_ptr;
@@ -65,28 +69,16 @@ BSONObj toBson(const intrusive_ptr<DocumentSource>& source) {
     return arr[0].getDocument().toBson();
 }
 
-class CollectionBase {
+class DocumentSourceCursorTest : public unittest::Test {
 public:
-    CollectionBase() : client(&_opCtx) {}
-
-    ~CollectionBase() {
-        client.dropCollection(nss.ns());
+    DocumentSourceCursorTest()
+        : client(_opCtx.get()),
+          _ctx(new ExpressionContextForTest(_opCtx.get(), AggregationRequest(nss, {}))) {
+        _ctx->tempDir = storageGlobalParams.dbpath + "/_tmp";
     }
 
-protected:
-    const ServiceContext::UniqueOperationContext _opCtxPtr = cc().makeOperationContext();
-    OperationContext& _opCtx = *_opCtxPtr;
-    DBDirectClient client;
-};
-
-namespace DocumentSourceCursor {
-
-using mongo::DocumentSourceCursor;
-
-class Base : public CollectionBase {
-public:
-    Base() : _ctx(new ExpressionContextForTest(&_opCtx, AggregationRequest(nss, {}))) {
-        _ctx->tempDir = storageGlobalParams.dbpath + "/_tmp";
+    virtual ~DocumentSourceCursorTest() {
+        client.dropCollection(nss.ns());
     }
 
 protected:
@@ -94,16 +86,16 @@ protected:
         // clean up first if this was called before
         _source.reset();
 
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+        OldClientWriteContext ctx(opCtx(), nss.ns());
 
         auto qr = stdx::make_unique<QueryRequest>(nss);
         if (hint) {
             qr->setHint(*hint);
         }
-        auto cq = uassertStatusOK(CanonicalQuery::canonicalize(&_opCtx, std::move(qr)));
+        auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
 
         auto exec = uassertStatusOK(
-            getExecutor(&_opCtx, ctx.getCollection(), std::move(cq), PlanExecutor::NO_YIELD));
+            getExecutor(opCtx(), ctx.getCollection(), std::move(cq), PlanExecutor::NO_YIELD));
 
         exec->saveState();
         _source = DocumentSourceCursor::create(ctx.getCollection(), std::move(exec), _ctx);
@@ -117,6 +109,14 @@ protected:
         return _source.get();
     }
 
+    OperationContext* opCtx() {
+        return _opCtx.get();
+    }
+
+protected:
+    const ServiceContext::UniqueOperationContext _opCtx = cc().makeOperationContext();
+    DBDirectClient client;
+
 private:
     // It is important that these are ordered to ensure correct destruction order.
     intrusive_ptr<ExpressionContextForTest> _ctx;
@@ -124,78 +124,66 @@ private:
 };
 
 /** Create a DocumentSourceCursor. */
-class Empty : public Base {
-public:
-    void run() {
-        createSource();
-        // The DocumentSourceCursor doesn't hold a read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        // The collection is empty, so the source produces no results.
-        ASSERT(source()->getNext().isEOF());
-        // Exhausting the source releases the read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-    }
-};
+TEST_F(DocumentSourceCursorTest, Empty) {
+    createSource();
+    // The DocumentSourceCursor doesn't hold a read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    // The collection is empty, so the source produces no results.
+    ASSERT(source()->getNext().isEOF());
+    // Exhausting the source releases the read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+}
 
 /** Iterate a DocumentSourceCursor. */
-class Iterate : public Base {
-public:
-    void run() {
-        client.insert(nss.ns(), BSON("a" << 1));
-        createSource();
-        // The DocumentSourceCursor doesn't hold a read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        // The cursor will produce the expected result.
-        auto next = source()->getNext();
-        ASSERT(next.isAdvanced());
-        ASSERT_VALUE_EQ(Value(1), next.getDocument().getField("a"));
-        // There are no more results.
-        ASSERT(source()->getNext().isEOF());
-        // Exhausting the source releases the read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-    }
-};
+TEST_F(DocumentSourceCursorTest, Iterate) {
+    client.insert(nss.ns(), BSON("a" << 1));
+    createSource();
+    // The DocumentSourceCursor doesn't hold a read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    // The cursor will produce the expected result.
+    auto next = source()->getNext();
+    ASSERT(next.isAdvanced());
+    ASSERT_VALUE_EQ(Value(1), next.getDocument().getField("a"));
+    // There are no more results.
+    ASSERT(source()->getNext().isEOF());
+    // Exhausting the source releases the read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+}
 
 /** Dispose of a DocumentSourceCursor. */
-class Dispose : public Base {
-public:
-    void run() {
-        createSource();
-        // The DocumentSourceCursor doesn't hold a read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        source()->dispose();
-        // Releasing the cursor releases the read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        // The source is marked as exhausted.
-        ASSERT(source()->getNext().isEOF());
-    }
-};
+TEST_F(DocumentSourceCursorTest, Dispose) {
+    createSource();
+    // The DocumentSourceCursor doesn't hold a read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    source()->dispose();
+    // Releasing the cursor releases the read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    // The source is marked as exhausted.
+    ASSERT(source()->getNext().isEOF());
+}
 
 /** Iterate a DocumentSourceCursor and then dispose of it. */
-class IterateDispose : public Base {
-public:
-    void run() {
-        client.insert(nss.ns(), BSON("a" << 1));
-        client.insert(nss.ns(), BSON("a" << 2));
-        client.insert(nss.ns(), BSON("a" << 3));
-        createSource();
-        // The result is as expected.
-        auto next = source()->getNext();
-        ASSERT(next.isAdvanced());
-        ASSERT_VALUE_EQ(Value(1), next.getDocument().getField("a"));
-        // The next result is as expected.
-        next = source()->getNext();
-        ASSERT(next.isAdvanced());
-        ASSERT_VALUE_EQ(Value(2), next.getDocument().getField("a"));
-        // The DocumentSourceCursor doesn't hold a read lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        source()->dispose();
-        // Disposing of the source releases the lock.
-        ASSERT(!_opCtx.lockState()->isReadLocked());
-        // The source cannot be advanced further.
-        ASSERT(source()->getNext().isEOF());
-    }
-};
+TEST_F(DocumentSourceCursorTest, IterateDispose) {
+    client.insert(nss.ns(), BSON("a" << 1));
+    client.insert(nss.ns(), BSON("a" << 2));
+    client.insert(nss.ns(), BSON("a" << 3));
+    createSource();
+    // The result is as expected.
+    auto next = source()->getNext();
+    ASSERT(next.isAdvanced());
+    ASSERT_VALUE_EQ(Value(1), next.getDocument().getField("a"));
+    // The next result is as expected.
+    next = source()->getNext();
+    ASSERT(next.isAdvanced());
+    ASSERT_VALUE_EQ(Value(2), next.getDocument().getField("a"));
+    // The DocumentSourceCursor doesn't hold a read lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    source()->dispose();
+    // Disposing of the source releases the lock.
+    ASSERT(!opCtx()->lockState()->isReadLocked());
+    // The source cannot be advanced further.
+    ASSERT(source()->getNext().isEOF());
+}
 
 /** Set a value or await an expected value. */
 class PendingValue {
@@ -221,148 +209,267 @@ private:
 
 
 /** Test coalescing a limit into a cursor */
-class LimitCoalesce : public Base {
-public:
-    intrusive_ptr<DocumentSourceLimit> mkLimit(long long limit) {
-        return DocumentSourceLimit::create(ctx(), limit);
-    }
-    void run() {
-        client.insert(nss.ns(), BSON("a" << 1));
-        client.insert(nss.ns(), BSON("a" << 2));
-        client.insert(nss.ns(), BSON("a" << 3));
-        createSource();
-
-        Pipeline::SourceContainer container;
-        container.push_back(source());
-        container.push_back(mkLimit(10));
-        source()->optimizeAt(container.begin(), &container);
-
-        // initial limit becomes limit of cursor
-        ASSERT_EQUALS(container.size(), 1U);
-        ASSERT_EQUALS(source()->getLimit(), 10);
-
-        container.push_back(mkLimit(2));
-        source()->optimizeAt(container.begin(), &container);
-        // smaller limit lowers cursor limit
-        ASSERT_EQUALS(container.size(), 1U);
-        ASSERT_EQUALS(source()->getLimit(), 2);
-
-        container.push_back(mkLimit(3));
-        source()->optimizeAt(container.begin(), &container);
-        // higher limit doesn't effect cursor limit
-        ASSERT_EQUALS(container.size(), 1U);
-        ASSERT_EQUALS(source()->getLimit(), 2);
-
-        // The cursor allows exactly 2 documents through
-        ASSERT(source()->getNext().isAdvanced());
-        ASSERT(source()->getNext().isAdvanced());
-        ASSERT(source()->getNext().isEOF());
-    }
-};
+TEST_F(DocumentSourceCursorTest, LimitCoalesce) {
+    client.insert(nss.ns(), BSON("a" << 1));
+    client.insert(nss.ns(), BSON("a" << 2));
+    client.insert(nss.ns(), BSON("a" << 3));
+    createSource();
+
+    Pipeline::SourceContainer container;
+    container.push_back(source());
+    container.push_back(DocumentSourceLimit::create(ctx(), 10));
+    source()->optimizeAt(container.begin(), &container);
+
+    // initial limit becomes limit of cursor
+    ASSERT_EQUALS(container.size(), 1U);
+    ASSERT_EQUALS(source()->getLimit(), 10);
+
+    container.push_back(DocumentSourceLimit::create(ctx(), 2));
+    source()->optimizeAt(container.begin(), &container);
+    // smaller limit lowers cursor limit
+    ASSERT_EQUALS(container.size(), 1U);
+    ASSERT_EQUALS(source()->getLimit(), 2);
+
+    container.push_back(DocumentSourceLimit::create(ctx(), 3));
+    source()->optimizeAt(container.begin(), &container);
+    // higher limit doesn't effect cursor limit
+    ASSERT_EQUALS(container.size(), 1U);
+    ASSERT_EQUALS(source()->getLimit(), 2);
+
+    // The cursor allows exactly 2 documents through
+    ASSERT(source()->getNext().isAdvanced());
+    ASSERT(source()->getNext().isAdvanced());
+    ASSERT(source()->getNext().isEOF());
+}
 
 //
 // Test cursor output sort.
 //
-class CollectionScanProvidesNoSort : public Base {
-public:
-    void run() {
-        createSource(BSON("$natural" << 1));
-        ASSERT_EQ(source()->getOutputSorts().size(), 0U);
-        source()->dispose();
-    }
-};
-
-class IndexScanProvidesSortOnKeys : public Base {
-public:
-    void run() {
-        client.createIndex(nss.ns(), BSON("a" << 1));
-        createSource(BSON("a" << 1));
+TEST_F(DocumentSourceCursorTest, CollectionScanProvidesNoSort) {
+    createSource(BSON("$natural" << 1));
+    ASSERT_EQ(source()->getOutputSorts().size(), 0U);
+    source()->dispose();
+}
 
-        ASSERT_EQ(source()->getOutputSorts().size(), 1U);
-        ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1)), 1U);
-        source()->dispose();
-    }
-};
+TEST_F(DocumentSourceCursorTest, IndexScanProvidesSortOnKeys) {
+    client.createIndex(nss.ns(), BSON("a" << 1));
+    createSource(BSON("a" << 1));
 
-class ReverseIndexScanProvidesSort : public Base {
-public:
-    void run() {
-        client.createIndex(nss.ns(), BSON("a" << -1));
-        createSource(BSON("a" << -1));
+    ASSERT_EQ(source()->getOutputSorts().size(), 1U);
+    ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1)), 1U);
+    source()->dispose();
+}
 
-        ASSERT_EQ(source()->getOutputSorts().size(), 1U);
-        ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << -1)), 1U);
-        source()->dispose();
-    }
-};
+TEST_F(DocumentSourceCursorTest, ReverseIndexScanProvidesSort) {
+    client.createIndex(nss.ns(), BSON("a" << -1));
+    createSource(BSON("a" << -1));
 
-class CompoundIndexScanProvidesMultipleSorts : public Base {
-public:
-    void run() {
-        client.createIndex(nss.ns(), BSON("a" << 1 << "b" << -1));
-        createSource(BSON("a" << 1 << "b" << -1));
-
-        ASSERT_EQ(source()->getOutputSorts().size(), 2U);
-        ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1)), 1U);
-        ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1 << "b" << -1)), 1U);
-        source()->dispose();
-    }
-};
+    ASSERT_EQ(source()->getOutputSorts().size(), 1U);
+    ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << -1)), 1U);
+    source()->dispose();
+}
 
-class SerializationRespectsExplainModes : public Base {
-public:
-    void run() {
-        createSource();
+TEST_F(DocumentSourceCursorTest, CompoundIndexScanProvidesMultipleSorts) {
+    client.createIndex(nss.ns(), BSON("a" << 1 << "b" << -1));
+    createSource(BSON("a" << 1 << "b" << -1));
 
-        {
-            // Nothing serialized when no explain mode specified.
-            auto explainResult = source()->serialize();
-            ASSERT_TRUE(explainResult.missing());
-        }
+    ASSERT_EQ(source()->getOutputSorts().size(), 2U);
+    ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1)), 1U);
+    ASSERT_EQ(source()->getOutputSorts().count(BSON("a" << 1 << "b" << -1)), 1U);
+    source()->dispose();
+}
 
-        {
-            auto explainResult = source()->serialize(ExplainOptions::Verbosity::kQueryPlanner);
-            ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
-            ASSERT_TRUE(explainResult["$cursor"]["executionStats"].missing());
-        }
+TEST_F(DocumentSourceCursorTest, SerializationRespectsExplainModes) {
+    createSource();
 
-        {
-            auto explainResult = source()->serialize(ExplainOptions::Verbosity::kExecStats);
-            ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
-            ASSERT_FALSE(explainResult["$cursor"]["executionStats"].missing());
-            ASSERT_TRUE(explainResult["$cursor"]["executionStats"]["allPlansExecution"].missing());
-        }
+    {
+        // Nothing serialized when no explain mode specified.
+        auto explainResult = source()->serialize();
+        ASSERT_TRUE(explainResult.missing());
+    }
 
-        {
-            auto explainResult =
-                source()->serialize(ExplainOptions::Verbosity::kExecAllPlans).getDocument();
-            ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
-            ASSERT_FALSE(explainResult["$cursor"]["executionStats"].missing());
-            ASSERT_FALSE(explainResult["$cursor"]["executionStats"]["allPlansExecution"].missing());
-        }
-        source()->dispose();
+    {
+        auto explainResult = source()->serialize(ExplainOptions::Verbosity::kQueryPlanner);
+        ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
+        ASSERT_TRUE(explainResult["$cursor"]["executionStats"].missing());
     }
-};
 
-}  // namespace DocumentSourceCursor
+    {
+        auto explainResult = source()->serialize(ExplainOptions::Verbosity::kExecStats);
+        ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
+        ASSERT_FALSE(explainResult["$cursor"]["executionStats"].missing());
+        ASSERT_TRUE(explainResult["$cursor"]["executionStats"]["allPlansExecution"].missing());
+    }
 
-class All : public Suite {
-public:
-    All() : Suite("documentsource") {}
-    void setupTests() {
-        add<DocumentSourceCursor::Empty>();
-        add<DocumentSourceCursor::Iterate>();
-        add<DocumentSourceCursor::Dispose>();
-        add<DocumentSourceCursor::IterateDispose>();
-        add<DocumentSourceCursor::LimitCoalesce>();
-        add<DocumentSourceCursor::CollectionScanProvidesNoSort>();
-        add<DocumentSourceCursor::IndexScanProvidesSortOnKeys>();
-        add<DocumentSourceCursor::ReverseIndexScanProvidesSort>();
-        add<DocumentSourceCursor::CompoundIndexScanProvidesMultipleSorts>();
-        add<DocumentSourceCursor::SerializationRespectsExplainModes>();
+    {
+        auto explainResult =
+            source()->serialize(ExplainOptions::Verbosity::kExecAllPlans).getDocument();
+        ASSERT_FALSE(explainResult["$cursor"]["queryPlanner"].missing());
+        ASSERT_FALSE(explainResult["$cursor"]["executionStats"].missing());
+        ASSERT_FALSE(explainResult["$cursor"]["executionStats"]["allPlansExecution"].missing());
     }
-};
+    source()->dispose();
+}
 
-SuiteInstance<All> myall;
+TEST_F(DocumentSourceCursorTest, TailableAwaitDataCursorStillUsableAfterTimeout) {
+    // Make sure the collection exists, otherwise we'll default to a NO_YIELD yield policy.
+    const bool capped = true;
+    const bool cappedSize = 1024;
+    ASSERT_TRUE(client.createCollection(nss.ns(), cappedSize, capped));
+    client.insert(nss.ns(), BSON("a" << 1));
+
+    // Make a tailable collection scan wrapped up in a PlanExecutor.
+    AutoGetCollectionForRead readLock(opCtx(), nss);
+    auto workingSet = stdx::make_unique<WorkingSet>();
+    CollectionScanParams collScanParams;
+    collScanParams.collection = readLock.getCollection();
+    collScanParams.tailable = true;
+    auto filter = BSON("a" << 1);
+    auto matchExpression =
+        uassertStatusOK(MatchExpressionParser::parse(filter, ctx()->getCollator()));
+    auto collectionScan = stdx::make_unique<CollectionScan>(
+        opCtx(), collScanParams, workingSet.get(), matchExpression.get());
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(filter);
+    queryRequest->setTailable(true);
+    queryRequest->setAwaitData(true);
+    auto canonicalQuery = unittest::assertGet(
+        CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest), nullptr));
+    auto planExecutor =
+        uassertStatusOK(PlanExecutor::make(opCtx(),
+                                           std::move(workingSet),
+                                           std::move(collectionScan),
+                                           std::move(canonicalQuery),
+                                           readLock.getCollection(),
+                                           PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT));
+
+    // Make a DocumentSourceCursor.
+    ctx()->tailableMode = ExpressionContext::TailableMode::kTailableAndAwaitData;
+    // DocumentSourceCursor expects a PlanExecutor that has had its state saved.
+    planExecutor->saveState();
+    auto cursor =
+        DocumentSourceCursor::create(readLock.getCollection(), std::move(planExecutor), ctx());
+
+    ASSERT(cursor->getNext().isEOF());
+    cursor->dispose();
+}
+
+TEST_F(DocumentSourceCursorTest, NonAwaitDataCursorShouldErrorAfterTimeout) {
+    // Make sure the collection exists, otherwise we'll default to a NO_YIELD yield policy.
+    ASSERT_TRUE(client.createCollection(nss.ns()));
+    client.insert(nss.ns(), BSON("a" << 1));
+
+    // Make a tailable collection scan wrapped up in a PlanExecutor.
+    AutoGetCollectionForRead readLock(opCtx(), nss);
+    auto workingSet = stdx::make_unique<WorkingSet>();
+    CollectionScanParams collScanParams;
+    collScanParams.collection = readLock.getCollection();
+    auto filter = BSON("a" << 1);
+    auto matchExpression =
+        uassertStatusOK(MatchExpressionParser::parse(filter, ctx()->getCollator()));
+    auto collectionScan = stdx::make_unique<CollectionScan>(
+        opCtx(), collScanParams, workingSet.get(), matchExpression.get());
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(filter);
+    auto canonicalQuery = unittest::assertGet(
+        CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest), nullptr));
+    auto planExecutor =
+        uassertStatusOK(PlanExecutor::make(opCtx(),
+                                           std::move(workingSet),
+                                           std::move(collectionScan),
+                                           std::move(canonicalQuery),
+                                           readLock.getCollection(),
+                                           PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT));
+
+    // Make a DocumentSourceCursor.
+    ctx()->tailableMode = ExpressionContext::TailableMode::kNormal;
+    // DocumentSourceCursor expects a PlanExecutor that has had its state saved.
+    planExecutor->saveState();
+    auto cursor =
+        DocumentSourceCursor::create(readLock.getCollection(), std::move(planExecutor), ctx());
+
+    ON_BLOCK_EXIT([cursor]() { cursor->dispose(); });
+    ASSERT_THROWS_CODE(cursor->getNext().isEOF(), AssertionException, ErrorCodes::QueryPlanKilled);
+}
+
+TEST_F(DocumentSourceCursorTest, TailableAwaitDataCursorShouldErrorAfterBeingKilled) {
+    // Make sure the collection exists, otherwise we'll default to a NO_YIELD yield policy.
+    const bool capped = true;
+    const bool cappedSize = 1024;
+    ASSERT_TRUE(client.createCollection(nss.ns(), cappedSize, capped));
+    client.insert(nss.ns(), BSON("a" << 1));
+
+    // Make a tailable collection scan wrapped up in a PlanExecutor.
+    AutoGetCollectionForRead readLock(opCtx(), nss);
+    auto workingSet = stdx::make_unique<WorkingSet>();
+    CollectionScanParams collScanParams;
+    collScanParams.collection = readLock.getCollection();
+    collScanParams.tailable = true;
+    auto filter = BSON("a" << 1);
+    auto matchExpression =
+        uassertStatusOK(MatchExpressionParser::parse(filter, ctx()->getCollator()));
+    auto collectionScan = stdx::make_unique<CollectionScan>(
+        opCtx(), collScanParams, workingSet.get(), matchExpression.get());
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(filter);
+    auto canonicalQuery = unittest::assertGet(
+        CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest), nullptr));
+    auto planExecutor =
+        uassertStatusOK(PlanExecutor::make(opCtx(),
+                                           std::move(workingSet),
+                                           std::move(collectionScan),
+                                           std::move(canonicalQuery),
+                                           readLock.getCollection(),
+                                           PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED));
+
+    // Make a DocumentSourceCursor.
+    ctx()->tailableMode = ExpressionContext::TailableMode::kTailableAndAwaitData;
+    // DocumentSourceCursor expects a PlanExecutor that has had its state saved.
+    planExecutor->saveState();
+    auto cursor =
+        DocumentSourceCursor::create(readLock.getCollection(), std::move(planExecutor), ctx());
+
+    ON_BLOCK_EXIT([cursor]() { cursor->dispose(); });
+    ASSERT_THROWS_CODE(cursor->getNext().isEOF(), AssertionException, ErrorCodes::QueryPlanKilled);
+}
+
+TEST_F(DocumentSourceCursorTest, NormalCursorShouldErrorAfterBeingKilled) {
+    // Make sure the collection exists, otherwise we'll default to a NO_YIELD yield policy.
+    ASSERT_TRUE(client.createCollection(nss.ns()));
+    client.insert(nss.ns(), BSON("a" << 1));
+
+    // Make a tailable collection scan wrapped up in a PlanExecutor.
+    AutoGetCollectionForRead readLock(opCtx(), nss);
+    auto workingSet = stdx::make_unique<WorkingSet>();
+    CollectionScanParams collScanParams;
+    collScanParams.collection = readLock.getCollection();
+    auto filter = BSON("a" << 1);
+    auto matchExpression =
+        uassertStatusOK(MatchExpressionParser::parse(filter, ctx()->getCollator()));
+    auto collectionScan = stdx::make_unique<CollectionScan>(
+        opCtx(), collScanParams, workingSet.get(), matchExpression.get());
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(filter);
+    auto canonicalQuery = unittest::assertGet(
+        CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest), nullptr));
+    auto planExecutor =
+        uassertStatusOK(PlanExecutor::make(opCtx(),
+                                           std::move(workingSet),
+                                           std::move(collectionScan),
+                                           std::move(canonicalQuery),
+                                           readLock.getCollection(),
+                                           PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED));
+
+    // Make a DocumentSourceCursor.
+    ctx()->tailableMode = ExpressionContext::TailableMode::kNormal;
+    // DocumentSourceCursor expects a PlanExecutor that has had its state saved.
+    planExecutor->saveState();
+    auto cursor =
+        DocumentSourceCursor::create(readLock.getCollection(), std::move(planExecutor), ctx());
+
+    ON_BLOCK_EXIT([cursor]() { cursor->dispose(); });
+    ASSERT_THROWS_CODE(cursor->getNext().isEOF(), AssertionException, ErrorCodes::QueryPlanKilled);
+}
 
-}  // namespace DocumentSourceCursorTests
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/dbtests/executor_registry.cpp b/src/mongo/dbtests/executor_registry.cpp
index 8d961c51eff..044499c4416 100644
--- a/src/mongo/dbtests/executor_registry.cpp
+++ b/src/mongo/dbtests/executor_registry.cpp
@@ -137,7 +137,7 @@ public:
         // At this point, we're done yielding.  We recover our lock.
 
         // And clean up anything that happened before.
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         // Make sure that the PlanExecutor moved forward over the deleted data.  We don't see
         // foo==10
@@ -164,14 +164,12 @@ public:
             ASSERT_EQUALS(i, obj["foo"].numberInt());
         }
 
-        // Save state and register.
         exec->saveState();
 
         // Drop a collection that's not ours.
         _client.dropCollection("unittests.someboguscollection");
 
-        // Unregister and restore state.
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&obj, NULL));
         ASSERT_EQUALS(10, obj["foo"].numberInt());
@@ -180,10 +178,7 @@ public:
 
         _client.dropCollection(nss.ns());
 
-        exec->restoreState();
-
-        // PlanExecutor was killed.
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&obj, NULL));
+        ASSERT_EQUALS(ErrorCodes::QueryPlanKilled, exec->restoreState());
     }
 };
 
@@ -204,8 +199,7 @@ public:
 
         exec->saveState();
         _client.dropIndexes(nss.ns());
-        exec->restoreState();
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&obj, NULL));
+        ASSERT_EQUALS(ErrorCodes::QueryPlanKilled, exec->restoreState());
     }
 };
 
@@ -226,8 +220,7 @@ public:
 
         exec->saveState();
         _client.dropIndex(nss.ns(), BSON("foo" << 1));
-        exec->restoreState();
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&obj, NULL));
+        ASSERT_EQUALS(ErrorCodes::QueryPlanKilled, exec->restoreState());
     }
 };
 
@@ -251,7 +244,7 @@ public:
         _ctx.reset();
         _client.dropDatabase("somesillydb");
         _ctx.reset(new OldClientWriteContext(&_opCtx, nss.ns()));
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&obj, NULL));
         ASSERT_EQUALS(10, obj["foo"].numberInt());
@@ -262,11 +255,7 @@ public:
         _ctx.reset();
         _client.dropDatabase("unittests");
         _ctx.reset(new OldClientWriteContext(&_opCtx, nss.ns()));
-        exec->restoreState();
-        _ctx.reset();
-
-        // PlanExecutor was killed.
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&obj, NULL));
+        ASSERT_EQUALS(ErrorCodes::QueryPlanKilled, exec->restoreState());
     }
 };
 
diff --git a/src/mongo/dbtests/query_plan_executor.cpp b/src/mongo/dbtests/query_plan_executor.cpp
index 4b1a1eef0e4..99f4548b631 100644
--- a/src/mongo/dbtests/query_plan_executor.cpp
+++ b/src/mongo/dbtests/query_plan_executor.cpp
@@ -53,7 +53,8 @@
 #include "mongo/dbtests/dbtests.h"
 #include "mongo/stdx/memory.h"
 
-namespace QueryPlanExecutor {
+namespace mongo {
+namespace {
 
 using std::shared_ptr;
 using std::string;
@@ -62,11 +63,11 @@ using stdx::make_unique;
 
 static const NamespaceString nss("unittests.QueryPlanExecutor");
 
-class PlanExecutorBase {
+class PlanExecutorTest : public unittest::Test {
 public:
-    PlanExecutorBase() : _client(&_opCtx) {}
+    PlanExecutorTest() : _client(&_opCtx) {}
 
-    virtual ~PlanExecutorBase() {
+    virtual ~PlanExecutorTest() {
         _client.dropCollection(nss.ns());
     }
 
@@ -94,8 +95,12 @@ public:
      * Given a match expression, represented as the BSON object 'filterObj', create a PlanExecutor
      * capable of executing a simple collection scan.
      */
-    unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeCollScanExec(Collection* coll,
-                                                                     BSONObj& filterObj) {
+    unique_ptr<PlanExecutor, PlanExecutor::Deleter> makeCollScanExec(
+        Collection* coll,
+        BSONObj& filterObj,
+        PlanExecutor::YieldPolicy yieldPolicy = PlanExecutor::YieldPolicy::YIELD_MANUAL,
+        bool tailable = false,
+        bool awaitData = false) {
         CollectionScanParams csparams;
         csparams.collection = coll;
         csparams.direction = CollectionScanParams::FORWARD;
@@ -104,8 +109,10 @@ public:
         // Canonicalize the query.
         auto qr = stdx::make_unique<QueryRequest>(nss);
         qr->setFilter(filterObj);
+        qr->setTailable(tailable);
+        qr->setAwaitData(awaitData);
         auto statusWithCQ = CanonicalQuery::canonicalize(&_opCtx, std::move(qr));
-        verify(statusWithCQ.isOK());
+        ASSERT_OK(statusWithCQ.getStatus());
         unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
         verify(NULL != cq.get());
 
@@ -114,12 +121,8 @@ public:
             new CollectionScan(&_opCtx, csparams, ws.get(), cq.get()->root()));
 
         // Hand the plan off to the executor.
-        auto statusWithPlanExecutor = PlanExecutor::make(&_opCtx,
-                                                         std::move(ws),
-                                                         std::move(root),
-                                                         std::move(cq),
-                                                         coll,
-                                                         PlanExecutor::YIELD_MANUAL);
+        auto statusWithPlanExecutor = PlanExecutor::make(
+            &_opCtx, std::move(ws), std::move(root), std::move(cq), coll, yieldPolicy);
         ASSERT_OK(statusWithPlanExecutor.getStatus());
         return std::move(statusWithPlanExecutor.getValue());
     }
@@ -191,105 +194,160 @@ private:
  * Test dropping the collection while the
  * PlanExecutor is doing a collection scan.
  */
-class DropCollScan : public PlanExecutorBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        insert(BSON("_id" << 1));
-        insert(BSON("_id" << 2));
+TEST_F(PlanExecutorTest, DropCollScan) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1));
+    insert(BSON("_id" << 2));
 
-        BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
+    BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
 
-        Collection* coll = ctx.getCollection();
-        auto exec = makeCollScanExec(coll, filterObj);
+    Collection* coll = ctx.getCollection();
+    auto exec = makeCollScanExec(coll, filterObj);
 
-        BSONObj objOut;
-        ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
-        ASSERT_EQUALS(1, objOut["_id"].numberInt());
+    BSONObj objOut;
+    ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
+    ASSERT_EQUALS(1, objOut["_id"].numberInt());
 
-        // After dropping the collection, the plan executor should be dead.
-        dropCollection();
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&objOut, NULL));
-    }
-};
+    // After dropping the collection, the plan executor should be dead.
+    dropCollection();
+    ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&objOut, NULL));
+}
 
 /**
  * Test dropping the collection while the PlanExecutor is doing an index scan.
  */
-class DropIndexScan : public PlanExecutorBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        insert(BSON("_id" << 1 << "a" << 6));
-        insert(BSON("_id" << 2 << "a" << 7));
-        insert(BSON("_id" << 3 << "a" << 8));
-        BSONObj indexSpec = BSON("a" << 1);
-        addIndex(indexSpec);
+TEST_F(PlanExecutorTest, DropIndexScan) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1 << "a" << 6));
+    insert(BSON("_id" << 2 << "a" << 7));
+    insert(BSON("_id" << 3 << "a" << 8));
+    BSONObj indexSpec = BSON("a" << 1);
+    addIndex(indexSpec);
 
-        auto exec = makeIndexScanExec(ctx.db(), indexSpec, 7, 10);
+    auto exec = makeIndexScanExec(ctx.db(), indexSpec, 7, 10);
 
-        BSONObj objOut;
-        ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
-        ASSERT_EQUALS(7, objOut["a"].numberInt());
+    BSONObj objOut;
+    ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
+    ASSERT_EQUALS(7, objOut["a"].numberInt());
 
-        // After dropping the collection, the plan executor should be dead.
-        dropCollection();
-        ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&objOut, NULL));
-    }
-};
+    // After dropping the collection, the plan executor should be dead.
+    dropCollection();
+    ASSERT_EQUALS(PlanExecutor::DEAD, exec->getNext(&objOut, NULL));
+}
 
 /**
  * Test dropping the collection while an agg PlanExecutor is doing an index scan.
  */
-class DropIndexScanAgg : public PlanExecutorBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-
-        insert(BSON("_id" << 1 << "a" << 6));
-        insert(BSON("_id" << 2 << "a" << 7));
-        insert(BSON("_id" << 3 << "a" << 8));
-        BSONObj indexSpec = BSON("a" << 1);
-        addIndex(indexSpec);
-
-        Collection* collection = ctx.getCollection();
-
-        // Create the aggregation pipeline.
-        std::vector<BSONObj> rawPipeline = {fromjson("{$match: {a: {$gte: 7, $lte: 10}}}")};
-        boost::intrusive_ptr<ExpressionContextForTest> expCtx =
-            new ExpressionContextForTest(&_opCtx, AggregationRequest(nss, rawPipeline));
-
-        // Create an "inner" plan executor and register it with the cursor manager so that it can
-        // get notified when the collection is dropped.
-        unique_ptr<PlanExecutor, PlanExecutor::Deleter> innerExec(
-            makeIndexScanExec(ctx.db(), indexSpec, 7, 10));
-
-        // Wrap the "inner" plan executor in a DocumentSourceCursor and add it as the first source
-        // in the pipeline.
-        innerExec->saveState();
-        auto cursorSource = DocumentSourceCursor::create(collection, std::move(innerExec), expCtx);
-        auto pipeline = assertGet(Pipeline::create({cursorSource}, expCtx));
+TEST_F(PlanExecutorTest, DropIndexScanAgg) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+
+    insert(BSON("_id" << 1 << "a" << 6));
+    insert(BSON("_id" << 2 << "a" << 7));
+    insert(BSON("_id" << 3 << "a" << 8));
+    BSONObj indexSpec = BSON("a" << 1);
+    addIndex(indexSpec);
+
+    Collection* collection = ctx.getCollection();
+
+    // Create the aggregation pipeline.
+    std::vector<BSONObj> rawPipeline = {fromjson("{$match: {a: {$gte: 7, $lte: 10}}}")};
+    boost::intrusive_ptr<ExpressionContextForTest> expCtx =
+        new ExpressionContextForTest(&_opCtx, AggregationRequest(nss, rawPipeline));
+
+    // Create an "inner" plan executor and register it with the cursor manager so that it can
+    // get notified when the collection is dropped.
+    unique_ptr<PlanExecutor, PlanExecutor::Deleter> innerExec(
+        makeIndexScanExec(ctx.db(), indexSpec, 7, 10));
+
+    // Wrap the "inner" plan executor in a DocumentSourceCursor and add it as the first source
+    // in the pipeline.
+    innerExec->saveState();
+    auto cursorSource = DocumentSourceCursor::create(collection, std::move(innerExec), expCtx);
+    auto pipeline = assertGet(Pipeline::create({cursorSource}, expCtx));
+
+    // Create the output PlanExecutor that pulls results from the pipeline.
+    auto ws = make_unique<WorkingSet>();
+    auto proxy = make_unique<PipelineProxyStage>(&_opCtx, std::move(pipeline), ws.get());
+
+    auto statusWithPlanExecutor = PlanExecutor::make(
+        &_opCtx, std::move(ws), std::move(proxy), collection, PlanExecutor::NO_YIELD);
+    ASSERT_OK(statusWithPlanExecutor.getStatus());
+    auto outerExec = std::move(statusWithPlanExecutor.getValue());
+
+    dropCollection();
+
+    // Verify that the aggregation pipeline returns an error because its "inner" plan executor
+    // has been killed due to the collection being dropped.
+    BSONObj objOut;
+    ASSERT_THROWS_CODE(
+        outerExec->getNext(&objOut, nullptr), AssertionException, ErrorCodes::QueryPlanKilled);
+}
+
+TEST_F(PlanExecutorTest, ShouldReportErrorIfExceedsTimeLimitDuringYield) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1));
+    insert(BSON("_id" << 2));
+
+    BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
+
+    Collection* coll = ctx.getCollection();
+    auto exec = makeCollScanExec(coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT);
+
+    BSONObj resultObj;
+    ASSERT_EQ(PlanExecutor::DEAD, exec->getNext(&resultObj, nullptr));
+    ASSERT_EQ(ErrorCodes::ExceededTimeLimit, WorkingSetCommon::getMemberObjectStatus(resultObj));
+}
+
+TEST_F(PlanExecutorTest, ShouldReportEOFIfExceedsTimeLimitDuringYieldButIsTailableAndAwaitData) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1));
+    insert(BSON("_id" << 2));
+
+    BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
+
+    Collection* coll = ctx.getCollection();
+    const bool tailable = true;
+    const bool awaitData = true;
+    auto exec = makeCollScanExec(
+        coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT, tailable, awaitData);
+
+    BSONObj resultObj;
+    ASSERT_EQ(PlanExecutor::IS_EOF, exec->getNext(&resultObj, nullptr));
+}
+
+TEST_F(PlanExecutorTest, ShouldNotSwallowExceedsTimeLimitDuringYieldButIsTailableButNotAwaitData) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1));
+    insert(BSON("_id" << 2));
+
+    BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
+
+    Collection* coll = ctx.getCollection();
+    const bool tailable = true;
+    auto exec =
+        makeCollScanExec(coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_TIME_OUT, tailable);
+
+    BSONObj resultObj;
+    ASSERT_EQ(PlanExecutor::DEAD, exec->getNext(&resultObj, nullptr));
+    ASSERT_EQ(ErrorCodes::ExceededTimeLimit, WorkingSetCommon::getMemberObjectStatus(resultObj));
+}
+
+TEST_F(PlanExecutorTest, ShouldReportErrorIfKilledDuringYield) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    insert(BSON("_id" << 1));
+    insert(BSON("_id" << 2));
+
+    BSONObj filterObj = fromjson("{_id: {$gt: 0}}");
+
+    Collection* coll = ctx.getCollection();
+    auto exec = makeCollScanExec(coll, filterObj, PlanExecutor::YieldPolicy::ALWAYS_MARK_KILLED);
+
+    BSONObj resultObj;
+    ASSERT_EQ(PlanExecutor::DEAD, exec->getNext(&resultObj, nullptr));
+    ASSERT_EQ(ErrorCodes::QueryPlanKilled, WorkingSetCommon::getMemberObjectStatus(resultObj));
+}
 
-        // Create the output PlanExecutor that pulls results from the pipeline.
-        auto ws = make_unique<WorkingSet>();
-        auto proxy = make_unique<PipelineProxyStage>(&_opCtx, std::move(pipeline), ws.get());
-
-        auto statusWithPlanExecutor = PlanExecutor::make(
-            &_opCtx, std::move(ws), std::move(proxy), collection, PlanExecutor::NO_YIELD);
-        ASSERT_OK(statusWithPlanExecutor.getStatus());
-        auto outerExec = std::move(statusWithPlanExecutor.getValue());
-
-        dropCollection();
-
-        // Verify that the aggregation pipeline returns an error because its "inner" plan executor
-        // has been killed due to the collection being dropped.
-        BSONObj objOut;
-        ASSERT_THROWS_CODE(
-            outerExec->getNext(&objOut, nullptr), AssertionException, ErrorCodes::QueryPlanKilled);
-    }
-};
-
-class SnapshotBase : public PlanExecutorBase {
+class PlanExecutorSnapshotTest : public PlanExecutorTest {
 protected:
     void setupCollection() {
         insert(BSON("_id" << 1 << "a" << 1));
@@ -338,70 +396,50 @@ protected:
  * twice due to a concurrent document move and collection
  * scan.
  */
-class SnapshotControl : public SnapshotBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        setupCollection();
+TEST_F(PlanExecutorSnapshotTest, SnapshotControl) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    setupCollection();
 
-        BSONObj filterObj = fromjson("{a: {$gte: 2}}");
+    BSONObj filterObj = fromjson("{a: {$gte: 2}}");
 
-        Collection* coll = ctx.getCollection();
-        auto exec = makeCollScanExec(coll, filterObj);
+    Collection* coll = ctx.getCollection();
+    auto exec = makeCollScanExec(coll, filterObj);
 
-        BSONObj objOut;
-        ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
-        ASSERT_EQUALS(2, objOut["a"].numberInt());
+    BSONObj objOut;
+    ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
+    ASSERT_EQUALS(2, objOut["a"].numberInt());
 
-        forceDocumentMove();
+    forceDocumentMove();
 
-        int ids[] = {3, 4, 2};
-        checkIds(ids, exec.get());
-    }
-};
+    int ids[] = {3, 4, 2};
+    checkIds(ids, exec.get());
+}
 
 /**
  * A snapshot is really just a hint that means scan the _id index.
  * Make sure that we do not see the document move with an _id
  * index scan.
  */
-class SnapshotTest : public SnapshotBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        setupCollection();
-        BSONObj indexSpec = BSON("_id" << 1);
-        addIndex(indexSpec);
+TEST_F(PlanExecutorSnapshotTest, SnapshotTest) {
+    OldClientWriteContext ctx(&_opCtx, nss.ns());
+    setupCollection();
+    BSONObj indexSpec = BSON("_id" << 1);
+    addIndex(indexSpec);
 
-        BSONObj filterObj = fromjson("{a: {$gte: 2}}");
-        auto exec = makeIndexScanExec(ctx.db(), indexSpec, 2, 5);
+    BSONObj filterObj = fromjson("{a: {$gte: 2}}");
+    auto exec = makeIndexScanExec(ctx.db(), indexSpec, 2, 5);
 
-        BSONObj objOut;
-        ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
-        ASSERT_EQUALS(2, objOut["a"].numberInt());
-
-        forceDocumentMove();
+    BSONObj objOut;
+    ASSERT_EQUALS(PlanExecutor::ADVANCED, exec->getNext(&objOut, NULL));
+    ASSERT_EQUALS(2, objOut["a"].numberInt());
 
-        // Since this time we're scanning the _id index,
-        // we should not see the moved document again.
-        int ids[] = {3, 4};
-        checkIds(ids, exec.get());
-    }
-};
-
-class All : public Suite {
-public:
-    All() : Suite("query_plan_executor") {}
-
-    void setupTests() {
-        add<DropCollScan>();
-        add<DropIndexScan>();
-        add<DropIndexScanAgg>();
-        add<SnapshotControl>();
-        add<SnapshotTest>();
-    }
-};
+    forceDocumentMove();
 
-SuiteInstance<All> queryPlanExecutorAll;
+    // Since this time we're scanning the _id index,
+    // we should not see the moved document again.
+    int ids[] = {3, 4};
+    checkIds(ids, exec.get());
+}
 
-}  // namespace QueryPlanExecutor
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/dbtests/query_stage_multiplan.cpp b/src/mongo/dbtests/query_stage_multiplan.cpp
index 0f02e120110..dd1fc50b924 100644
--- a/src/mongo/dbtests/query_stage_multiplan.cpp
+++ b/src/mongo/dbtests/query_stage_multiplan.cpp
@@ -44,6 +44,7 @@
 #include "mongo/db/matcher/expression_parser.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/query/get_executor.h"
+#include "mongo/db/query/mock_yield_policies.h"
 #include "mongo/db/query/plan_executor.h"
 #include "mongo/db/query/plan_summary_stats.h"
 #include "mongo/db/query/query_knobs.h"
@@ -61,9 +62,7 @@ const std::unique_ptr<ClockSource> clockSource = stdx::make_unique<ClockSourceMo
 // How we access the external setParameter testing bool.
 extern AtomicBool internalQueryForceIntersectionPlans;
 
-}  // namespace mongo
-
-namespace QueryStageMultiPlan {
+namespace {
 
 using std::unique_ptr;
 using std::vector;
@@ -82,40 +81,43 @@ QuerySolution* createQuerySolution() {
     return soln.release();
 }
 
-class QueryStageMultiPlanBase {
+class QueryStageMultiPlanTest : public unittest::Test {
 public:
-    QueryStageMultiPlanBase() : _client(&_opCtx) {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+    QueryStageMultiPlanTest() : _client(_opCtx.get()) {
+        OldClientWriteContext ctx(_opCtx.get(), nss.ns());
         _client.dropCollection(nss.ns());
     }
 
-    virtual ~QueryStageMultiPlanBase() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+    virtual ~QueryStageMultiPlanTest() {
+        OldClientWriteContext ctx(_opCtx.get(), nss.ns());
         _client.dropCollection(nss.ns());
     }
 
     void addIndex(const BSONObj& obj) {
-        ASSERT_OK(dbtests::createIndex(&_opCtx, nss.ns(), obj));
+        ASSERT_OK(dbtests::createIndex(_opCtx.get(), nss.ns(), obj));
     }
 
     void insert(const BSONObj& obj) {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+        OldClientWriteContext ctx(_opCtx.get(), nss.ns());
         _client.insert(nss.ns(), obj);
     }
 
     void remove(const BSONObj& obj) {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+        OldClientWriteContext ctx(_opCtx.get(), nss.ns());
         _client.remove(nss.ns(), obj);
     }
 
     OperationContext* opCtx() {
-        return &_opCtx;
+        return _opCtx.get();
+    }
+
+    ServiceContext* serviceContext() {
+        return _opCtx->getServiceContext();
     }
 
 protected:
-    const ServiceContext::UniqueOperationContext _txnPtr = cc().makeOperationContext();
-    OperationContext& _opCtx = *_txnPtr;
-    ClockSource* const _clock = _opCtx.getServiceContext()->getFastClockSource();
+    const ServiceContext::UniqueOperationContext _opCtx = cc().makeOperationContext();
+    ClockSource* const _clock = _opCtx->getServiceContext()->getFastClockSource();
 
     DBDirectClient _client;
 };
@@ -123,326 +125,424 @@ protected:
 
 // Basic ranking test: collection scan vs. highly selective index scan.  Make sure we also get
 // all expected results out as well.
-class MPSCollectionScanVsHighlySelectiveIXScan : public QueryStageMultiPlanBase {
-public:
-    void run() {
-        const int N = 5000;
-        for (int i = 0; i < N; ++i) {
-            insert(BSON("foo" << (i % 10)));
-        }
+TEST_F(QueryStageMultiPlanTest, MPSCollectionScanVsHighlySelectiveIXScan) {
+    const int N = 5000;
+    for (int i = 0; i < N; ++i) {
+        insert(BSON("foo" << (i % 10)));
+    }
 
-        addIndex(BSON("foo" << 1));
-
-        AutoGetCollectionForReadCommand ctx(&_opCtx, nss);
-        const Collection* coll = ctx.getCollection();
-
-        // Plan 0: IXScan over foo == 7
-        // Every call to work() returns something so this should clearly win (by current scoring
-        // at least).
-        std::vector<IndexDescriptor*> indexes;
-        coll->getIndexCatalog()->findIndexesByKeyPattern(
-            &_opCtx, BSON("foo" << 1), false, &indexes);
-        ASSERT_EQ(indexes.size(), 1U);
-
-        IndexScanParams ixparams;
-        ixparams.descriptor = indexes[0];
-        ixparams.bounds.isSimpleRange = true;
-        ixparams.bounds.startKey = BSON("" << 7);
-        ixparams.bounds.endKey = BSON("" << 7);
-        ixparams.bounds.boundInclusion = BoundInclusion::kIncludeBothStartAndEndKeys;
-        ixparams.direction = 1;
-
-        unique_ptr<WorkingSet> sharedWs(new WorkingSet());
-        IndexScan* ix = new IndexScan(&_opCtx, ixparams, sharedWs.get(), NULL);
-        unique_ptr<PlanStage> firstRoot(new FetchStage(&_opCtx, sharedWs.get(), ix, NULL, coll));
-
-        // Plan 1: CollScan with matcher.
-        CollectionScanParams csparams;
-        csparams.collection = coll;
-        csparams.direction = CollectionScanParams::FORWARD;
-
-        // Make the filter.
-        BSONObj filterObj = BSON("foo" << 7);
-        const CollatorInterface* collator = nullptr;
-        StatusWithMatchExpression statusWithMatcher =
-            MatchExpressionParser::parse(filterObj, collator);
-        verify(statusWithMatcher.isOK());
-        unique_ptr<MatchExpression> filter = std::move(statusWithMatcher.getValue());
-        // Make the stage.
-        unique_ptr<PlanStage> secondRoot(
-            new CollectionScan(&_opCtx, csparams, sharedWs.get(), filter.get()));
-
-        // Hand the plans off to the MPS.
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(BSON("foo" << 7));
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        verify(statusWithCQ.isOK());
-        unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
-        verify(NULL != cq.get());
-
-        unique_ptr<MultiPlanStage> mps =
-            make_unique<MultiPlanStage>(&_opCtx, ctx.getCollection(), cq.get());
-        mps->addPlan(createQuerySolution(), firstRoot.release(), sharedWs.get());
-        mps->addPlan(createQuerySolution(), secondRoot.release(), sharedWs.get());
-
-        // Plan 0 aka the first plan aka the index scan should be the best.
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        mps->pickBestPlan(&yieldPolicy).transitional_ignore();
-        ASSERT(mps->bestPlanChosen());
-        ASSERT_EQUALS(0, mps->bestPlanIdx());
-
-        // Takes ownership of arguments other than 'collection'.
-        auto statusWithPlanExecutor = PlanExecutor::make(&_opCtx,
-                                                         std::move(sharedWs),
-                                                         std::move(mps),
-                                                         std::move(cq),
-                                                         coll,
-                                                         PlanExecutor::NO_YIELD);
-        ASSERT_OK(statusWithPlanExecutor.getStatus());
-        auto exec = std::move(statusWithPlanExecutor.getValue());
-
-        // Get all our results out.
-        int results = 0;
-        BSONObj obj;
-        PlanExecutor::ExecState state;
-        while (PlanExecutor::ADVANCED == (state = exec->getNext(&obj, NULL))) {
-            ASSERT_EQUALS(obj["foo"].numberInt(), 7);
-            ++results;
-        }
-        ASSERT_EQUALS(PlanExecutor::IS_EOF, state);
-        ASSERT_EQUALS(results, N / 10);
+    addIndex(BSON("foo" << 1));
+
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+    const Collection* coll = ctx.getCollection();
+
+    // Plan 0: IXScan over foo == 7
+    // Every call to work() returns something so this should clearly win (by current scoring
+    // at least).
+    std::vector<IndexDescriptor*> indexes;
+    coll->getIndexCatalog()->findIndexesByKeyPattern(
+        _opCtx.get(), BSON("foo" << 1), false, &indexes);
+    ASSERT_EQ(indexes.size(), 1U);
+
+    IndexScanParams ixparams;
+    ixparams.descriptor = indexes[0];
+    ixparams.bounds.isSimpleRange = true;
+    ixparams.bounds.startKey = BSON("" << 7);
+    ixparams.bounds.endKey = BSON("" << 7);
+    ixparams.bounds.boundInclusion = BoundInclusion::kIncludeBothStartAndEndKeys;
+    ixparams.direction = 1;
+
+    unique_ptr<WorkingSet> sharedWs(new WorkingSet());
+    IndexScan* ix = new IndexScan(_opCtx.get(), ixparams, sharedWs.get(), NULL);
+    unique_ptr<PlanStage> firstRoot(new FetchStage(_opCtx.get(), sharedWs.get(), ix, NULL, coll));
+
+    // Plan 1: CollScan with matcher.
+    CollectionScanParams csparams;
+    csparams.collection = coll;
+    csparams.direction = CollectionScanParams::FORWARD;
+
+    // Make the filter.
+    BSONObj filterObj = BSON("foo" << 7);
+    const CollatorInterface* collator = nullptr;
+    StatusWithMatchExpression statusWithMatcher = MatchExpressionParser::parse(filterObj, collator);
+    verify(statusWithMatcher.isOK());
+    unique_ptr<MatchExpression> filter = std::move(statusWithMatcher.getValue());
+    // Make the stage.
+    unique_ptr<PlanStage> secondRoot(
+        new CollectionScan(_opCtx.get(), csparams, sharedWs.get(), filter.get()));
+
+    // Hand the plans off to the MPS.
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(BSON("foo" << 7));
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    verify(statusWithCQ.isOK());
+    unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+    verify(NULL != cq.get());
+
+    unique_ptr<MultiPlanStage> mps =
+        make_unique<MultiPlanStage>(_opCtx.get(), ctx.getCollection(), cq.get());
+    mps->addPlan(createQuerySolution(), firstRoot.release(), sharedWs.get());
+    mps->addPlan(createQuerySolution(), secondRoot.release(), sharedWs.get());
+
+    // Plan 0 aka the first plan aka the index scan should be the best.
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(mps->pickBestPlan(&yieldPolicy));
+    ASSERT(mps->bestPlanChosen());
+    ASSERT_EQUALS(0, mps->bestPlanIdx());
+
+    // Takes ownership of arguments other than 'collection'.
+    auto statusWithPlanExecutor = PlanExecutor::make(_opCtx.get(),
+                                                     std::move(sharedWs),
+                                                     std::move(mps),
+                                                     std::move(cq),
+                                                     coll,
+                                                     PlanExecutor::NO_YIELD);
+    ASSERT_OK(statusWithPlanExecutor.getStatus());
+    auto exec = std::move(statusWithPlanExecutor.getValue());
+
+    // Get all our results out.
+    int results = 0;
+    BSONObj obj;
+    PlanExecutor::ExecState state;
+    while (PlanExecutor::ADVANCED == (state = exec->getNext(&obj, NULL))) {
+        ASSERT_EQUALS(obj["foo"].numberInt(), 7);
+        ++results;
     }
-};
+    ASSERT_EQUALS(PlanExecutor::IS_EOF, state);
+    ASSERT_EQUALS(results, N / 10);
+}
 
 // Case in which we select a blocking plan as the winner, and a non-blocking plan
 // is available as a backup.
-class MPSBackupPlan : public QueryStageMultiPlanBase {
-public:
-    void run() {
-        // Data is just a single {_id: 1, a: 1, b: 1} document.
-        insert(BSON("_id" << 1 << "a" << 1 << "b" << 1));
-
-        // Indices on 'a' and 'b'.
-        addIndex(BSON("a" << 1));
-        addIndex(BSON("b" << 1));
-
-        AutoGetCollectionForReadCommand ctx(&_opCtx, nss);
-        Collection* collection = ctx.getCollection();
-
-        // Query for both 'a' and 'b' and sort on 'b'.
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(BSON("a" << 1 << "b" << 1));
-        qr->setSort(BSON("b" << 1));
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        verify(statusWithCQ.isOK());
-        unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
-        ASSERT(NULL != cq.get());
-
-        // Force index intersection.
-        bool forceIxisectOldValue = internalQueryForceIntersectionPlans.load();
-        internalQueryForceIntersectionPlans.store(true);
-
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
-        // Turn this off otherwise it pops up in some plans.
-        plannerParams.options &= ~QueryPlannerParams::KEEP_MUTATIONS;
-
-        // Plan.
-        vector<QuerySolution*> solutions;
-        Status status = QueryPlanner::plan(*cq, plannerParams, &solutions);
-        ASSERT(status.isOK());
-
-        // We expect a plan using index {a: 1} and plan using index {b: 1} and
-        // an index intersection plan.
-        ASSERT_EQUALS(solutions.size(), 3U);
-
-        // Fill out the MultiPlanStage.
-        unique_ptr<MultiPlanStage> mps(new MultiPlanStage(&_opCtx, collection, cq.get()));
-        unique_ptr<WorkingSet> ws(new WorkingSet());
-        // Put each solution from the planner into the MPR.
-        for (size_t i = 0; i < solutions.size(); ++i) {
-            PlanStage* root;
-            ASSERT(StageBuilder::build(&_opCtx, collection, *cq, *solutions[i], ws.get(), &root));
-            // Takes ownership of 'solutions[i]' and 'root'.
-            mps->addPlan(solutions[i], root, ws.get());
-        }
+TEST_F(QueryStageMultiPlanTest, MPSBackupPlan) {
+    // Data is just a single {_id: 1, a: 1, b: 1} document.
+    insert(BSON("_id" << 1 << "a" << 1 << "b" << 1));
+
+    // Indices on 'a' and 'b'.
+    addIndex(BSON("a" << 1));
+    addIndex(BSON("b" << 1));
+
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+    Collection* collection = ctx.getCollection();
+
+    // Query for both 'a' and 'b' and sort on 'b'.
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(BSON("a" << 1 << "b" << 1));
+    qr->setSort(BSON("b" << 1));
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    verify(statusWithCQ.isOK());
+    unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+    ASSERT(NULL != cq.get());
+
+    // Force index intersection.
+    bool forceIxisectOldValue = internalQueryForceIntersectionPlans.load();
+    internalQueryForceIntersectionPlans.store(true);
+
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(_opCtx.get(), collection, cq.get(), &plannerParams);
+    // Turn this off otherwise it pops up in some plans.
+    plannerParams.options &= ~QueryPlannerParams::KEEP_MUTATIONS;
+
+    // Plan.
+    vector<QuerySolution*> solutions;
+    Status status = QueryPlanner::plan(*cq, plannerParams, &solutions);
+    ASSERT(status.isOK());
+
+    // We expect a plan using index {a: 1} and plan using index {b: 1} and
+    // an index intersection plan.
+    ASSERT_EQUALS(solutions.size(), 3U);
+
+    // Fill out the MultiPlanStage.
+    unique_ptr<MultiPlanStage> mps(new MultiPlanStage(_opCtx.get(), collection, cq.get()));
+    unique_ptr<WorkingSet> ws(new WorkingSet());
+    // Put each solution from the planner into the MPR.
+    for (size_t i = 0; i < solutions.size(); ++i) {
+        PlanStage* root;
+        ASSERT(StageBuilder::build(_opCtx.get(), collection, *cq, *solutions[i], ws.get(), &root));
+        // Takes ownership of 'solutions[i]' and 'root'.
+        mps->addPlan(solutions[i], root, ws.get());
+    }
 
-        // This sets a backup plan.
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        mps->pickBestPlan(&yieldPolicy).transitional_ignore();
-        ASSERT(mps->bestPlanChosen());
-        ASSERT(mps->hasBackupPlan());
-
-        // We should have picked the index intersection plan due to forcing ixisect.
-        QuerySolution* soln = mps->bestSolution();
-        ASSERT(QueryPlannerTestLib::solutionMatches(
-            "{sort: {pattern: {b: 1}, limit: 0, node: {sortKeyGen: {node:"
-            "{fetch: {node: {andSorted: {nodes: ["
-            "{ixscan: {filter: null, pattern: {a:1}}},"
-            "{ixscan: {filter: null, pattern: {b:1}}}]}}}}}}}}",
-            soln->root.get()));
-
-        // Get the resulting document.
-        PlanStage::StageState state = PlanStage::NEED_TIME;
-        WorkingSetID wsid;
-        while (state != PlanStage::ADVANCED) {
-            state = mps->work(&wsid);
-        }
-        WorkingSetMember* member = ws->get(wsid);
-
-        // Check the document returned by the query.
-        ASSERT(member->hasObj());
-        BSONObj expectedDoc = BSON("_id" << 1 << "a" << 1 << "b" << 1);
-        ASSERT(expectedDoc.woCompare(member->obj.value()) == 0);
-
-        // The blocking plan became unblocked, so we should no longer have a backup plan,
-        // and the winning plan should still be the index intersection one.
-        ASSERT(!mps->hasBackupPlan());
-        soln = mps->bestSolution();
-        ASSERT(QueryPlannerTestLib::solutionMatches(
-            "{sort: {pattern: {b: 1}, limit: 0, node: {sortKeyGen: {node:"
-            "{fetch: {node: {andSorted: {nodes: ["
-            "{ixscan: {filter: null, pattern: {a:1}}},"
-            "{ixscan: {filter: null, pattern: {b:1}}}]}}}}}}}}",
-            soln->root.get()));
-
-        // Restore index intersection force parameter.
-        internalQueryForceIntersectionPlans.store(forceIxisectOldValue);
+    // This sets a backup plan.
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(mps->pickBestPlan(&yieldPolicy));
+    ASSERT(mps->bestPlanChosen());
+    ASSERT(mps->hasBackupPlan());
+
+    // We should have picked the index intersection plan due to forcing ixisect.
+    QuerySolution* soln = mps->bestSolution();
+    ASSERT(QueryPlannerTestLib::solutionMatches(
+        "{sort: {pattern: {b: 1}, limit: 0, node: {sortKeyGen: {node:"
+        "{fetch: {node: {andSorted: {nodes: ["
+        "{ixscan: {filter: null, pattern: {a:1}}},"
+        "{ixscan: {filter: null, pattern: {b:1}}}]}}}}}}}}",
+        soln->root.get()));
+
+    // Get the resulting document.
+    PlanStage::StageState state = PlanStage::NEED_TIME;
+    WorkingSetID wsid;
+    while (state != PlanStage::ADVANCED) {
+        state = mps->work(&wsid);
     }
-};
+    WorkingSetMember* member = ws->get(wsid);
+
+    // Check the document returned by the query.
+    ASSERT(member->hasObj());
+    BSONObj expectedDoc = BSON("_id" << 1 << "a" << 1 << "b" << 1);
+    ASSERT(expectedDoc.woCompare(member->obj.value()) == 0);
+
+    // The blocking plan became unblocked, so we should no longer have a backup plan,
+    // and the winning plan should still be the index intersection one.
+    ASSERT(!mps->hasBackupPlan());
+    soln = mps->bestSolution();
+    ASSERT(QueryPlannerTestLib::solutionMatches(
+        "{sort: {pattern: {b: 1}, limit: 0, node: {sortKeyGen: {node:"
+        "{fetch: {node: {andSorted: {nodes: ["
+        "{ixscan: {filter: null, pattern: {a:1}}},"
+        "{ixscan: {filter: null, pattern: {b:1}}}]}}}}}}}}",
+        soln->root.get()));
+
+    // Restore index intersection force parameter.
+    internalQueryForceIntersectionPlans.store(forceIxisectOldValue);
+}
 
-// Test the structure and values of the explain output.
-class MPSExplainAllPlans : public QueryStageMultiPlanBase {
-public:
-    void run() {
-        // Insert a document to create the collection.
-        insert(BSON("x" << 1));
+/**
+ * Allocates a new WorkingSetMember with data 'dataObj' in 'ws', and adds the WorkingSetMember
+ * to 'qds'.
+ */
+void addMember(QueuedDataStage* qds, WorkingSet* ws, BSONObj dataObj) {
+    WorkingSetID id = ws->allocate();
+    WorkingSetMember* wsm = ws->get(id);
+    wsm->obj = Snapshotted<BSONObj>(SnapshotId(), BSON("x" << 1));
+    wsm->transitionToOwnedObj();
+    qds->pushBack(id);
+}
 
-        const int nDocs = 500;
+// Test the structure and values of the explain output.
+TEST_F(QueryStageMultiPlanTest, MPSExplainAllPlans) {
+    // Insert a document to create the collection.
+    insert(BSON("x" << 1));
 
-        auto ws = stdx::make_unique<WorkingSet>();
-        auto firstPlan = stdx::make_unique<QueuedDataStage>(&_opCtx, ws.get());
-        auto secondPlan = stdx::make_unique<QueuedDataStage>(&_opCtx, ws.get());
+    const int nDocs = 500;
 
-        for (int i = 0; i < nDocs; ++i) {
-            addMember(firstPlan.get(), ws.get(), BSON("x" << 1));
+    auto ws = stdx::make_unique<WorkingSet>();
+    auto firstPlan = stdx::make_unique<QueuedDataStage>(_opCtx.get(), ws.get());
+    auto secondPlan = stdx::make_unique<QueuedDataStage>(_opCtx.get(), ws.get());
 
-            // Make the second plan slower by inserting a NEED_TIME between every result.
-            addMember(secondPlan.get(), ws.get(), BSON("x" << 1));
-            secondPlan->pushBack(PlanStage::NEED_TIME);
-        }
+    for (int i = 0; i < nDocs; ++i) {
+        addMember(firstPlan.get(), ws.get(), BSON("x" << 1));
 
-        AutoGetCollectionForReadCommand ctx(&_opCtx, nss);
-
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(BSON("x" << 1));
-        auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
-        unique_ptr<MultiPlanStage> mps =
-            make_unique<MultiPlanStage>(&_opCtx, ctx.getCollection(), cq.get());
-
-        // Put each plan into the MultiPlanStage. Takes ownership of 'firstPlan' and 'secondPlan'.
-        auto firstSoln = stdx::make_unique<QuerySolution>();
-        auto secondSoln = stdx::make_unique<QuerySolution>();
-        mps->addPlan(firstSoln.release(), firstPlan.release(), ws.get());
-        mps->addPlan(secondSoln.release(), secondPlan.release(), ws.get());
-
-        // Making a PlanExecutor chooses the best plan.
-        auto exec = uassertStatusOK(PlanExecutor::make(
-            &_opCtx, std::move(ws), std::move(mps), ctx.getCollection(), PlanExecutor::NO_YIELD));
-
-        auto root = static_cast<MultiPlanStage*>(exec->getRootStage());
-        ASSERT_TRUE(root->bestPlanChosen());
-        // The first QueuedDataStage should have won.
-        ASSERT_EQ(root->bestPlanIdx(), 0);
-
-        BSONObjBuilder bob;
-        Explain::explainStages(
-            exec.get(), ctx.getCollection(), ExplainOptions::Verbosity::kExecAllPlans, &bob);
-        BSONObj explained = bob.done();
-
-        ASSERT_EQ(explained["executionStats"]["nReturned"].Int(), nDocs);
-        ASSERT_EQ(explained["executionStats"]["executionStages"]["needTime"].Int(), 0);
-        auto allPlansStats = explained["executionStats"]["allPlansExecution"].Array();
-        ASSERT_EQ(allPlansStats.size(), 2UL);
-        for (auto&& planStats : allPlansStats) {
-            int maxEvaluationResults = internalQueryPlanEvaluationMaxResults.load();
-            ASSERT_EQ(planStats["executionStages"]["stage"].String(), "QUEUED_DATA");
-            if (planStats["executionStages"]["needTime"].Int() > 0) {
-                // This is the losing plan. Should only have advanced about half the time.
-                ASSERT_LT(planStats["nReturned"].Int(), maxEvaluationResults);
-            } else {
-                // This is the winning plan. Stats here should be from the trial period.
-                ASSERT_EQ(planStats["nReturned"].Int(), maxEvaluationResults);
-            }
-        }
+        // Make the second plan slower by inserting a NEED_TIME between every result.
+        addMember(secondPlan.get(), ws.get(), BSON("x" << 1));
+        secondPlan->pushBack(PlanStage::NEED_TIME);
     }
 
-private:
-    /**
-     * Allocates a new WorkingSetMember with data 'dataObj' in 'ws', and adds the WorkingSetMember
-     * to 'qds'.
-     */
-    void addMember(QueuedDataStage* qds, WorkingSet* ws, BSONObj dataObj) {
-        WorkingSetID id = ws->allocate();
-        WorkingSetMember* wsm = ws->get(id);
-        wsm->obj = Snapshotted<BSONObj>(SnapshotId(), BSON("x" << 1));
-        wsm->transitionToOwnedObj();
-        qds->pushBack(id);
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(BSON("x" << 1));
+    auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
+    unique_ptr<MultiPlanStage> mps =
+        make_unique<MultiPlanStage>(_opCtx.get(), ctx.getCollection(), cq.get());
+
+    // Put each plan into the MultiPlanStage. Takes ownership of 'firstPlan' and 'secondPlan'.
+    auto firstSoln = stdx::make_unique<QuerySolution>();
+    auto secondSoln = stdx::make_unique<QuerySolution>();
+    mps->addPlan(firstSoln.release(), firstPlan.release(), ws.get());
+    mps->addPlan(secondSoln.release(), secondPlan.release(), ws.get());
+
+    // Making a PlanExecutor chooses the best plan.
+    auto exec = uassertStatusOK(PlanExecutor::make(
+        _opCtx.get(), std::move(ws), std::move(mps), ctx.getCollection(), PlanExecutor::NO_YIELD));
+
+    auto root = static_cast<MultiPlanStage*>(exec->getRootStage());
+    ASSERT_TRUE(root->bestPlanChosen());
+    // The first QueuedDataStage should have won.
+    ASSERT_EQ(root->bestPlanIdx(), 0);
+
+    BSONObjBuilder bob;
+    Explain::explainStages(
+        exec.get(), ctx.getCollection(), ExplainOptions::Verbosity::kExecAllPlans, &bob);
+    BSONObj explained = bob.done();
+
+    ASSERT_EQ(explained["executionStats"]["nReturned"].Int(), nDocs);
+    ASSERT_EQ(explained["executionStats"]["executionStages"]["needTime"].Int(), 0);
+    auto allPlansStats = explained["executionStats"]["allPlansExecution"].Array();
+    ASSERT_EQ(allPlansStats.size(), 2UL);
+    for (auto&& planStats : allPlansStats) {
+        int maxEvaluationResults = internalQueryPlanEvaluationMaxResults.load();
+        ASSERT_EQ(planStats["executionStages"]["stage"].String(), "QUEUED_DATA");
+        if (planStats["executionStages"]["needTime"].Int() > 0) {
+            // This is the losing plan. Should only have advanced about half the time.
+            ASSERT_LT(planStats["nReturned"].Int(), maxEvaluationResults);
+        } else {
+            // This is the winning plan. Stats here should be from the trial period.
+            ASSERT_EQ(planStats["nReturned"].Int(), maxEvaluationResults);
+        }
     }
-};
+}
 
 // Test that the plan summary only includes stats from the winning plan.
 //
 // This is a regression test for SERVER-20111.
-class MPSSummaryStats : public QueryStageMultiPlanBase {
-public:
-    void run() {
-        const int N = 5000;
-        for (int i = 0; i < N; ++i) {
-            insert(BSON("foo" << (i % 10)));
-        }
+TEST_F(QueryStageMultiPlanTest, MPSSummaryStats) {
+    const int N = 5000;
+    for (int i = 0; i < N; ++i) {
+        insert(BSON("foo" << (i % 10)));
+    }
 
-        // Add two indices to give more plans.
-        addIndex(BSON("foo" << 1));
-        addIndex(BSON("foo" << -1 << "bar" << 1));
+    // Add two indices to give more plans.
+    addIndex(BSON("foo" << 1));
+    addIndex(BSON("foo" << -1 << "bar" << 1));
 
-        AutoGetCollectionForReadCommand ctx(&_opCtx, nss);
-        Collection* coll = ctx.getCollection();
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+    Collection* coll = ctx.getCollection();
 
-        // Create the executor (Matching all documents).
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(BSON("foo" << BSON("$gte" << 0)));
-        auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
-        auto exec =
-            uassertStatusOK(getExecutor(&_opCtx, coll, std::move(cq), PlanExecutor::NO_YIELD));
+    // Create the executor (Matching all documents).
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(BSON("foo" << BSON("$gte" << 0)));
+    auto cq = uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
+    auto exec = uassertStatusOK(getExecutor(opCtx(), coll, std::move(cq), PlanExecutor::NO_YIELD));
+    ASSERT_EQ(exec->getRootStage()->stageType(), STAGE_MULTI_PLAN);
 
-        ASSERT_EQ(exec->getRootStage()->stageType(), STAGE_MULTI_PLAN);
+    ASSERT_OK(exec->executePlan());
 
-        exec->executePlan().transitional_ignore();
+    PlanSummaryStats stats;
+    Explain::getSummaryStats(*exec, &stats);
 
-        PlanSummaryStats stats;
-        Explain::getSummaryStats(*exec, &stats);
+    // If only the winning plan's stats are recorded, we should not have examined more than the
+    // total number of documents/index keys.
+    ASSERT_LTE(stats.totalDocsExamined, static_cast<size_t>(N));
+    ASSERT_LTE(stats.totalKeysExamined, static_cast<size_t>(N));
+}
 
-        // If only the winning plan's stats are recorded, we should not have examined more than the
-        // total number of documents/index keys.
-        ASSERT_LTE(stats.totalDocsExamined, static_cast<size_t>(N));
-        ASSERT_LTE(stats.totalKeysExamined, static_cast<size_t>(N));
+TEST_F(QueryStageMultiPlanTest, ShouldReportErrorIfExceedsTimeLimitDuringPlanning) {
+    const int N = 5000;
+    for (int i = 0; i < N; ++i) {
+        insert(BSON("foo" << (i % 10)));
     }
-};
 
-class All : public Suite {
-public:
-    All() : Suite("query_stage_multiplan") {}
+    // Add two indices to give more plans.
+    addIndex(BSON("foo" << 1));
+    addIndex(BSON("foo" << -1 << "bar" << 1));
+
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+    const auto coll = ctx.getCollection();
+
+    // Plan 0: IXScan over foo == 7
+    // Every call to work() returns something so this should clearly win (by current scoring
+    // at least).
+    std::vector<IndexDescriptor*> indexes;
+    coll->getIndexCatalog()->findIndexesByKeyPattern(
+        _opCtx.get(), BSON("foo" << 1), false, &indexes);
+    ASSERT_EQ(indexes.size(), 1U);
+
+    IndexScanParams ixparams;
+    ixparams.descriptor = indexes[0];
+    ixparams.bounds.isSimpleRange = true;
+    ixparams.bounds.startKey = BSON("" << 7);
+    ixparams.bounds.endKey = BSON("" << 7);
+    ixparams.bounds.boundInclusion = BoundInclusion::kIncludeBothStartAndEndKeys;
+    ixparams.direction = 1;
+
+    unique_ptr<WorkingSet> sharedWs(new WorkingSet());
+    IndexScan* ix = new IndexScan(_opCtx.get(), ixparams, sharedWs.get(), NULL);
+    unique_ptr<PlanStage> firstRoot(new FetchStage(_opCtx.get(), sharedWs.get(), ix, NULL, coll));
+
+    // Plan 1: CollScan with matcher.
+    CollectionScanParams csparams;
+    csparams.collection = coll;
+    csparams.direction = CollectionScanParams::FORWARD;
+
+    // Make the filter.
+    BSONObj filterObj = BSON("foo" << 7);
+    const CollatorInterface* collator = nullptr;
+    StatusWithMatchExpression statusWithMatcher = MatchExpressionParser::parse(filterObj, collator);
+    verify(statusWithMatcher.isOK());
+    unique_ptr<MatchExpression> filter = std::move(statusWithMatcher.getValue());
+    // Make the stage.
+    unique_ptr<PlanStage> secondRoot(
+        new CollectionScan(_opCtx.get(), csparams, sharedWs.get(), filter.get()));
+
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(BSON("foo" << 7));
+    auto canonicalQuery =
+        uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest)));
+    MultiPlanStage multiPlanStage(opCtx(),
+                                  ctx.getCollection(),
+                                  canonicalQuery.get(),
+                                  MultiPlanStage::CachingMode::NeverCache);
+    multiPlanStage.addPlan(createQuerySolution(), firstRoot.release(), sharedWs.get());
+    multiPlanStage.addPlan(createQuerySolution(), secondRoot.release(), sharedWs.get());
+
+    AlwaysTimeOutYieldPolicy alwaysTimeOutPolicy(serviceContext()->getFastClockSource());
+    ASSERT_EQ(ErrorCodes::ExceededTimeLimit, multiPlanStage.pickBestPlan(&alwaysTimeOutPolicy));
+}
 
-    void setupTests() {
-        add<MPSCollectionScanVsHighlySelectiveIXScan>();
-        add<MPSBackupPlan>();
-        add<MPSExplainAllPlans>();
-        add<MPSSummaryStats>();
+TEST_F(QueryStageMultiPlanTest, ShouldReportErrorIfKilledDuringPlanning) {
+    const int N = 5000;
+    for (int i = 0; i < N; ++i) {
+        insert(BSON("foo" << (i % 10)));
     }
-};
 
-SuiteInstance<All> queryStageMultiPlanAll;
+    // Add two indices to give more plans.
+    addIndex(BSON("foo" << 1));
+    addIndex(BSON("foo" << -1 << "bar" << 1));
+
+    AutoGetCollectionForReadCommand ctx(_opCtx.get(), nss);
+    const auto coll = ctx.getCollection();
+
+    // Plan 0: IXScan over foo == 7
+    // Every call to work() returns something so this should clearly win (by current scoring
+    // at least).
+    std::vector<IndexDescriptor*> indexes;
+    coll->getIndexCatalog()->findIndexesByKeyPattern(
+        _opCtx.get(), BSON("foo" << 1), false, &indexes);
+    ASSERT_EQ(indexes.size(), 1U);
+
+    IndexScanParams ixparams;
+    ixparams.descriptor = indexes[0];
+    ixparams.bounds.isSimpleRange = true;
+    ixparams.bounds.startKey = BSON("" << 7);
+    ixparams.bounds.endKey = BSON("" << 7);
+    ixparams.bounds.boundInclusion = BoundInclusion::kIncludeBothStartAndEndKeys;
+    ixparams.direction = 1;
+
+    unique_ptr<WorkingSet> sharedWs(new WorkingSet());
+    IndexScan* ix = new IndexScan(_opCtx.get(), ixparams, sharedWs.get(), NULL);
+    unique_ptr<PlanStage> firstRoot(new FetchStage(_opCtx.get(), sharedWs.get(), ix, NULL, coll));
+
+    // Plan 1: CollScan with matcher.
+    CollectionScanParams csparams;
+    csparams.collection = coll;
+    csparams.direction = CollectionScanParams::FORWARD;
+
+    // Make the filter.
+    BSONObj filterObj = BSON("foo" << 7);
+    const CollatorInterface* collator = nullptr;
+    StatusWithMatchExpression statusWithMatcher = MatchExpressionParser::parse(filterObj, collator);
+    verify(statusWithMatcher.isOK());
+    unique_ptr<MatchExpression> filter = std::move(statusWithMatcher.getValue());
+    // Make the stage.
+    unique_ptr<PlanStage> secondRoot(
+        new CollectionScan(_opCtx.get(), csparams, sharedWs.get(), filter.get()));
+
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(BSON("foo" << BSON("$gte" << 0)));
+    auto canonicalQuery =
+        uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest)));
+    MultiPlanStage multiPlanStage(opCtx(),
+                                  ctx.getCollection(),
+                                  canonicalQuery.get(),
+                                  MultiPlanStage::CachingMode::NeverCache);
+    multiPlanStage.addPlan(createQuerySolution(), firstRoot.release(), sharedWs.get());
+    multiPlanStage.addPlan(createQuerySolution(), secondRoot.release(), sharedWs.get());
+
+    AlwaysPlanKilledYieldPolicy alwaysPlanKilledYieldPolicy(serviceContext()->getFastClockSource());
+    ASSERT_EQ(ErrorCodes::QueryPlanKilled,
+              multiPlanStage.pickBestPlan(&alwaysPlanKilledYieldPolicy));
+}
 
-}  // namespace QueryStageMultiPlan
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/dbtests/query_stage_sort.cpp b/src/mongo/dbtests/query_stage_sort.cpp
index 2c698748b6c..6da81f409a7 100644
--- a/src/mongo/dbtests/query_stage_sort.cpp
+++ b/src/mongo/dbtests/query_stage_sort.cpp
@@ -365,7 +365,7 @@ public:
             coll->updateDocument(&_opCtx, *it, oldDoc, newDoc(oldDoc), false, false, NULL, &args);
             wuow.commit();
         }
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         // Read the rest of the data from the queued data stage.
         while (!queuedDataStage->isEOF()) {
@@ -385,7 +385,7 @@ public:
                 wuow.commit();
             }
         }
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         // Verify that it's sorted, the right number of documents are returned, and they're all
         // in the expected range.
@@ -465,7 +465,7 @@ public:
             coll->deleteDocument(&_opCtx, kUninitializedStmtId, *it++, nullOpDebug);
             wuow.commit();
         }
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         // Read the rest of the data from the queued data stage.
         while (!queuedDataStage->isEOF()) {
@@ -482,7 +482,7 @@ public:
                 wuow.commit();
             }
         }
-        exec->restoreState();
+        ASSERT_OK(exec->restoreState());
 
         // Regardless of storage engine, all the documents should come back with their objects
         int count = 0;
diff --git a/src/mongo/dbtests/query_stage_subplan.cpp b/src/mongo/dbtests/query_stage_subplan.cpp
index bca6559c25e..3b2e474e35a 100644
--- a/src/mongo/dbtests/query_stage_subplan.cpp
+++ b/src/mongo/dbtests/query_stage_subplan.cpp
@@ -32,6 +32,7 @@
 
 #include "mongo/db/catalog/collection.h"
 #include "mongo/db/catalog/database.h"
+#include "mongo/db/catalog/database_holder.h"
 #include "mongo/db/client.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbdirectclient.h"
@@ -42,23 +43,27 @@
 #include "mongo/db/pipeline/expression_context_for_test.h"
 #include "mongo/db/query/canonical_query.h"
 #include "mongo/db/query/get_executor.h"
+#include "mongo/db/query/mock_yield_policies.h"
+#include "mongo/db/query/query_test_service_context.h"
 #include "mongo/dbtests/dbtests.h"
+#include "mongo/util/assert_util.h"
 
-namespace QueryStageSubplan {
+namespace mongo {
+namespace {
 
 static const NamespaceString nss("unittests.QueryStageSubplan");
 
-class QueryStageSubplanBase {
+class QueryStageSubplanTest : public unittest::Test {
 public:
-    QueryStageSubplanBase() : _client(&_opCtx) {}
+    QueryStageSubplanTest() : _client(_opCtx.get()) {}
 
-    virtual ~QueryStageSubplanBase() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+    virtual ~QueryStageSubplanTest() {
+        OldClientWriteContext ctx(opCtx(), nss.ns());
         _client.dropCollection(nss.ns());
     }
 
     void addIndex(const BSONObj& obj) {
-        ASSERT_OK(dbtests::createIndex(&_opCtx, nss.ns(), obj));
+        ASSERT_OK(dbtests::createIndex(opCtx(), nss.ns(), obj));
     }
 
     void insert(const BSONObj& doc) {
@@ -66,7 +71,11 @@ public:
     }
 
     OperationContext* opCtx() {
-        return &_opCtx;
+        return _opCtx.get();
+    }
+
+    ServiceContext* serviceContext() {
+        return _opCtx->getServiceContext();
     }
 
 protected:
@@ -89,9 +98,8 @@ protected:
         return cq;
     }
 
-    const ServiceContext::UniqueOperationContext _txnPtr = cc().makeOperationContext();
-    OperationContext& _opCtx = *_txnPtr;
-    ClockSource* _clock = _opCtx.getServiceContext()->getFastClockSource();
+    const ServiceContext::UniqueOperationContext _opCtx = cc().makeOperationContext();
+    ClockSource* _clock = _opCtx->getServiceContext()->getFastClockSource();
 
 private:
     DBDirectClient _client;
@@ -103,542 +111,559 @@ private:
  * should gracefully fail after finding that no cache data is available, allowing us to fall
  * back to regular planning.
  */
-class QueryStageSubplanGeo2dOr : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        addIndex(BSON("a"
-                      << "2d"
-                      << "b"
-                      << 1));
-        addIndex(BSON("a"
-                      << "2d"));
-
-        BSONObj query = fromjson(
-            "{$or: [{a: {$geoWithin: {$centerSphere: [[0,0],10]}}},"
-            "{a: {$geoWithin: {$centerSphere: [[1,1],10]}}}]}");
-
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(query);
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        ASSERT_OK(statusWithCQ.getStatus());
-        std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
-
-        Collection* collection = ctx.getCollection();
-
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
-
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
-
-        // Plan selection should succeed due to falling back on regular planning.
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
-    }
-};
+TEST_F(QueryStageSubplanTest, QueryStageSubplanGeo2dOr) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
+    addIndex(BSON("a"
+                  << "2d"
+                  << "b"
+                  << 1));
+    addIndex(BSON("a"
+                  << "2d"));
+
+    BSONObj query = fromjson(
+        "{$or: [{a: {$geoWithin: {$centerSphere: [[0,0],10]}}},"
+        "{a: {$geoWithin: {$centerSphere: [[1,1],10]}}}]}");
+
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(query);
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    ASSERT_OK(statusWithCQ.getStatus());
+    std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+
+    Collection* collection = ctx.getCollection();
+
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
+
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
+
+    // Plan selection should succeed due to falling back on regular planning.
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+}
 
 /**
  * Test the SubplanStage's ability to plan an individual branch using the plan cache.
  */
-class QueryStageSubplanPlanFromCache : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+TEST_F(QueryStageSubplanTest, QueryStageSubplanPlanFromCache) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
 
-        addIndex(BSON("a" << 1));
-        addIndex(BSON("a" << 1 << "b" << 1));
-        addIndex(BSON("c" << 1));
+    addIndex(BSON("a" << 1));
+    addIndex(BSON("a" << 1 << "b" << 1));
+    addIndex(BSON("c" << 1));
 
-        for (int i = 0; i < 10; i++) {
-            insert(BSON("a" << 1 << "b" << i << "c" << i));
-        }
+    for (int i = 0; i < 10; i++) {
+        insert(BSON("a" << 1 << "b" << i << "c" << i));
+    }
 
-        // This query should result in a plan cache entry for the first $or branch, because
-        // there are two competing indices. The second branch has only one relevant index, so
-        // its winning plan should not be cached.
-        BSONObj query = fromjson("{$or: [{a: 1, b: 3}, {c: 1}]}");
+    // This query should result in a plan cache entry for the first $or branch, because
+    // there are two competing indices. The second branch has only one relevant index, so
+    // its winning plan should not be cached.
+    BSONObj query = fromjson("{$or: [{a: 1, b: 3}, {c: 1}]}");
 
-        Collection* collection = ctx.getCollection();
+    Collection* collection = ctx.getCollection();
 
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(query);
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        ASSERT_OK(statusWithCQ.getStatus());
-        std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(query);
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    ASSERT_OK(statusWithCQ.getStatus());
+    std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
 
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        // Nothing is in the cache yet, so neither branch should have been planned from
-        // the plan cache.
-        ASSERT_FALSE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+    // Nothing is in the cache yet, so neither branch should have been planned from
+    // the plan cache.
+    ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
 
-        // If we repeat the same query, the plan for the first branch should have come from
-        // the cache.
-        ws.clear();
-        subplan.reset(new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    // If we repeat the same query, the plan for the first branch should have come from
+    // the cache.
+    ws.clear();
+    subplan.reset(new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        ASSERT_TRUE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
-    }
-};
+    ASSERT_TRUE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+}
 
 /**
  * Ensure that the subplan stage doesn't create a plan cache entry if there are no query results.
  */
-class QueryStageSubplanDontCacheZeroResults : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+TEST_F(QueryStageSubplanTest, QueryStageSubplanDontCacheZeroResults) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
 
-        addIndex(BSON("a" << 1 << "b" << 1));
-        addIndex(BSON("a" << 1));
-        addIndex(BSON("c" << 1));
+    addIndex(BSON("a" << 1 << "b" << 1));
+    addIndex(BSON("a" << 1));
+    addIndex(BSON("c" << 1));
 
-        for (int i = 0; i < 10; i++) {
-            insert(BSON("a" << 1 << "b" << i << "c" << i));
-        }
+    for (int i = 0; i < 10; i++) {
+        insert(BSON("a" << 1 << "b" << i << "c" << i));
+    }
 
-        // Running this query should not create any cache entries. For the first branch, it's
-        // because there are no matching results. For the second branch it's because there is only
-        // one relevant index.
-        BSONObj query = fromjson("{$or: [{a: 1, b: 15}, {c: 1}]}");
+    // Running this query should not create any cache entries. For the first branch, it's
+    // because there are no matching results. For the second branch it's because there is only
+    // one relevant index.
+    BSONObj query = fromjson("{$or: [{a: 1, b: 15}, {c: 1}]}");
 
-        Collection* collection = ctx.getCollection();
+    Collection* collection = ctx.getCollection();
 
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(query);
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        ASSERT_OK(statusWithCQ.getStatus());
-        std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(query);
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    ASSERT_OK(statusWithCQ.getStatus());
+    std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
 
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        // Nothing is in the cache yet, so neither branch should have been planned from
-        // the plan cache.
-        ASSERT_FALSE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+    // Nothing is in the cache yet, so neither branch should have been planned from
+    // the plan cache.
+    ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
 
-        // If we run the query again, it should again be the case that neither branch gets planned
-        // from the cache (because the first call to pickBestPlan() refrained from creating any
-        // cache entries).
-        ws.clear();
-        subplan.reset(new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    // If we run the query again, it should again be the case that neither branch gets planned
+    // from the cache (because the first call to pickBestPlan() refrained from creating any
+    // cache entries).
+    ws.clear();
+    subplan.reset(new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        ASSERT_FALSE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
-    }
-};
+    ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+}
 
 /**
  * Ensure that the subplan stage doesn't create a plan cache entry if there are no query results.
  */
-class QueryStageSubplanDontCacheTies : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
+TEST_F(QueryStageSubplanTest, QueryStageSubplanDontCacheTies) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
 
-        addIndex(BSON("a" << 1 << "b" << 1));
-        addIndex(BSON("a" << 1 << "c" << 1));
-        addIndex(BSON("d" << 1));
+    addIndex(BSON("a" << 1 << "b" << 1));
+    addIndex(BSON("a" << 1 << "c" << 1));
+    addIndex(BSON("d" << 1));
 
-        for (int i = 0; i < 10; i++) {
-            insert(BSON("a" << 1 << "e" << 1 << "d" << 1));
-        }
+    for (int i = 0; i < 10; i++) {
+        insert(BSON("a" << 1 << "e" << 1 << "d" << 1));
+    }
 
-        // Running this query should not create any cache entries. For the first branch, it's
-        // because plans using the {a: 1, b: 1} and {a: 1, c: 1} indices should tie during plan
-        // ranking. For the second branch it's because there is only one relevant index.
-        BSONObj query = fromjson("{$or: [{a: 1, e: 1}, {d: 1}]}");
+    // Running this query should not create any cache entries. For the first branch, it's
+    // because plans using the {a: 1, b: 1} and {a: 1, c: 1} indices should tie during plan
+    // ranking. For the second branch it's because there is only one relevant index.
+    BSONObj query = fromjson("{$or: [{a: 1, e: 1}, {d: 1}]}");
 
-        Collection* collection = ctx.getCollection();
+    Collection* collection = ctx.getCollection();
 
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(query);
-        auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
-        ASSERT_OK(statusWithCQ.getStatus());
-        std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(query);
+    auto statusWithCQ = CanonicalQuery::canonicalize(opCtx(), std::move(qr));
+    ASSERT_OK(statusWithCQ.getStatus());
+    std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
 
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
 
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        // Nothing is in the cache yet, so neither branch should have been planned from
-        // the plan cache.
-        ASSERT_FALSE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+    // Nothing is in the cache yet, so neither branch should have been planned from
+    // the plan cache.
+    ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
 
-        // If we run the query again, it should again be the case that neither branch gets planned
-        // from the cache (because the first call to pickBestPlan() refrained from creating any
-        // cache entries).
-        ws.clear();
-        subplan.reset(new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
+    // If we run the query again, it should again be the case that neither branch gets planned
+    // from the cache (because the first call to pickBestPlan() refrained from creating any
+    // cache entries).
+    ws.clear();
+    subplan.reset(new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
 
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
 
-        ASSERT_FALSE(subplan->branchPlannedFromCache(0));
-        ASSERT_FALSE(subplan->branchPlannedFromCache(1));
-    }
-};
+    ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+    ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+}
 
 /**
  * Unit test the subplan stage's canUseSubplanning() method.
  */
-class QueryStageSubplanCanUseSubplanning : public QueryStageSubplanBase {
-public:
-    void run() {
-        // We won't try and subplan something that doesn't have an $or.
-        {
-            std::string findCmd = "{find: 'testns', filter: {$and:[{a:1}, {b:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+TEST_F(QueryStageSubplanTest, QueryStageSubplanCanUseSubplanning) {
+    // We won't try and subplan something that doesn't have an $or.
+    {
+        std::string findCmd = "{find: 'testns', filter: {$and:[{a:1}, {b:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Don't try and subplan if there is no filter.
-        {
-            std::string findCmd = "{find: 'testns'}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Don't try and subplan if there is no filter.
+    {
+        std::string findCmd = "{find: 'testns'}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // We won't try and subplan two contained ORs.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or:[{a:1}, {b:1}], $or:[{c:1}, {d:1}], e:1}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // We won't try and subplan two contained ORs.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or:[{a:1}, {b:1}], $or:[{c:1}, {d:1}], e:1}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning if there is a hint.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
-                "hint: {a:1, b:1}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning if there is a hint.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
+            "hint: {a:1, b:1}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning with min.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
-                "min: {a:1, b:1}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning with min.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
+            "min: {a:1, b:1}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning with max.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
-                "max: {a:2, b:2}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning with max.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
+            "max: {a:2, b:2}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning with tailable.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
-                "tailable: true}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning with tailable.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
+            "tailable: true}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning with snapshot.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
-                "snapshot: true}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning with snapshot.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]},"
+            "snapshot: true}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can use subplanning for rooted $or.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1, b:1}, {c:1, d:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_TRUE(SubplanStage::canUseSubplanning(*cq));
-
-            std::string findCmd2 =
-                "{find: 'testns',"
-                "filter: {$or: [{a:1}, {c:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq2 = cqFromFindCommand(findCmd2);
-            ASSERT_TRUE(SubplanStage::canUseSubplanning(*cq2));
-        }
+    // Can use subplanning for rooted $or.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1, b:1}, {c:1, d:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_TRUE(SubplanStage::canUseSubplanning(*cq));
+
+        std::string findCmd2 =
+            "{find: 'testns',"
+            "filter: {$or: [{a:1}, {c:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq2 = cqFromFindCommand(findCmd2);
+        ASSERT_TRUE(SubplanStage::canUseSubplanning(*cq2));
+    }
 
-        // Can't use subplanning for a single contained $or.
-        //
-        // TODO: Consider allowing this to use subplanning (see SERVER-13732).
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning for a single contained $or.
+    //
+    // TODO: Consider allowing this to use subplanning (see SERVER-13732).
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning if the contained $or query has a geo predicate.
-        //
-        // TODO: Consider allowing this to use subplanning (see SERVER-13732).
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {loc: {$geoWithin: {$centerSphere: [[0,0], 1]}},"
-                "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning if the contained $or query has a geo predicate.
+    //
+    // TODO: Consider allowing this to use subplanning (see SERVER-13732).
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {loc: {$geoWithin: {$centerSphere: [[0,0], 1]}},"
+            "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning if the contained $or query also has a $text predicate.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {$text: {$search: 'foo'},"
-                "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning if the contained $or query also has a $text predicate.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {$text: {$search: 'foo'},"
+            "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
+    }
 
-        // Can't use subplanning if the contained $or query also has a $near predicate.
-        {
-            std::string findCmd =
-                "{find: 'testns',"
-                "filter: {loc: {$near: [0, 0]},"
-                "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
-            std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
-            ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
-        }
+    // Can't use subplanning if the contained $or query also has a $near predicate.
+    {
+        std::string findCmd =
+            "{find: 'testns',"
+            "filter: {loc: {$near: [0, 0]},"
+            "e: 1, $or: [{a:1, b:1}, {c:1, d:1}]}}";
+        std::unique_ptr<CanonicalQuery> cq = cqFromFindCommand(findCmd);
+        ASSERT_FALSE(SubplanStage::canUseSubplanning(*cq));
     }
-};
+}
 
 /**
  * Unit test the subplan stage's rewriteToRootedOr() method.
  */
-class QueryStageSubplanRewriteToRootedOr : public QueryStageSubplanBase {
-public:
-    void run() {
-        // Rewrite (AND (OR a b) e) => (OR (AND a e) (AND b e))
-        {
-            BSONObj queryObj = fromjson("{$or:[{a:1}, {b:1}], e:1}");
-            const CollatorInterface* collator = nullptr;
-            StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
-            ASSERT_OK(expr.getStatus());
-            std::unique_ptr<MatchExpression> rewrittenExpr =
-                SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
-
-            std::string findCmdRewritten =
-                "{find: 'testns',"
-                "filter: {$or:[{a:1,e:1}, {b:1,e:1}]}}";
-            std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
-
-            ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
-        }
+TEST_F(QueryStageSubplanTest, QueryStageSubplanRewriteToRootedOr) {
+    // Rewrite (AND (OR a b) e) => (OR (AND a e) (AND b e))
+    {
+        BSONObj queryObj = fromjson("{$or:[{a:1}, {b:1}], e:1}");
+        const CollatorInterface* collator = nullptr;
+        StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
+        ASSERT_OK(expr.getStatus());
+        std::unique_ptr<MatchExpression> rewrittenExpr =
+            SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
+
+        std::string findCmdRewritten =
+            "{find: 'testns',"
+            "filter: {$or:[{a:1,e:1}, {b:1,e:1}]}}";
+        std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
+
+        ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
+    }
 
-        // Rewrite (AND (OR a b) e f) => (OR (AND a e f) (AND b e f))
-        {
-            BSONObj queryObj = fromjson("{$or:[{a:1}, {b:1}], e:1, f:1}");
-            const CollatorInterface* collator = nullptr;
-            StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
-            ASSERT_OK(expr.getStatus());
-            std::unique_ptr<MatchExpression> rewrittenExpr =
-                SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
-
-            std::string findCmdRewritten =
-                "{find: 'testns',"
-                "filter: {$or:[{a:1,e:1,f:1}, {b:1,e:1,f:1}]}}";
-            std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
-
-            ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
-        }
+    // Rewrite (AND (OR a b) e f) => (OR (AND a e f) (AND b e f))
+    {
+        BSONObj queryObj = fromjson("{$or:[{a:1}, {b:1}], e:1, f:1}");
+        const CollatorInterface* collator = nullptr;
+        StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
+        ASSERT_OK(expr.getStatus());
+        std::unique_ptr<MatchExpression> rewrittenExpr =
+            SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
+
+        std::string findCmdRewritten =
+            "{find: 'testns',"
+            "filter: {$or:[{a:1,e:1,f:1}, {b:1,e:1,f:1}]}}";
+        std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
+
+        ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
+    }
 
-        // Rewrite (AND (OR (AND a b) (AND c d) e f) => (OR (AND a b e f) (AND c d e f))
-        {
-            BSONObj queryObj = fromjson("{$or:[{a:1,b:1}, {c:1,d:1}], e:1,f:1}");
-            const CollatorInterface* collator = nullptr;
-            StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
-            ASSERT_OK(expr.getStatus());
-            std::unique_ptr<MatchExpression> rewrittenExpr =
-                SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
-
-            std::string findCmdRewritten =
-                "{find: 'testns',"
-                "filter: {$or:[{a:1,b:1,e:1,f:1},"
-                "{c:1,d:1,e:1,f:1}]}}";
-            std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
-
-            ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
-        }
+    // Rewrite (AND (OR (AND a b) (AND c d) e f) => (OR (AND a b e f) (AND c d e f))
+    {
+        BSONObj queryObj = fromjson("{$or:[{a:1,b:1}, {c:1,d:1}], e:1,f:1}");
+        const CollatorInterface* collator = nullptr;
+        StatusWithMatchExpression expr = MatchExpressionParser::parse(queryObj, collator);
+        ASSERT_OK(expr.getStatus());
+        std::unique_ptr<MatchExpression> rewrittenExpr =
+            SubplanStage::rewriteToRootedOr(std::move(expr.getValue()));
+
+        std::string findCmdRewritten =
+            "{find: 'testns',"
+            "filter: {$or:[{a:1,b:1,e:1,f:1},"
+            "{c:1,d:1,e:1,f:1}]}}";
+        std::unique_ptr<CanonicalQuery> cqRewritten = cqFromFindCommand(findCmdRewritten);
+
+        ASSERT(rewrittenExpr->equivalent(cqRewritten->root()));
     }
-};
+}
 
 /**
  * Test the subplan stage's ability to answer a contained $or query.
  */
-class QueryStageSubplanPlanContainedOr : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        addIndex(BSON("b" << 1 << "a" << 1));
-        addIndex(BSON("c" << 1 << "a" << 1));
-
-        BSONObj query = fromjson("{a: 1, $or: [{b: 2}, {c: 3}]}");
-
-        // Two of these documents match.
-        insert(BSON("_id" << 1 << "a" << 1 << "b" << 2));
-        insert(BSON("_id" << 2 << "a" << 2 << "b" << 2));
-        insert(BSON("_id" << 3 << "a" << 1 << "c" << 3));
-        insert(BSON("_id" << 4 << "a" << 1 << "c" << 4));
-
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(query);
-        auto cq = unittest::assertGet(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
-
-        Collection* collection = ctx.getCollection();
-
-        // Get planner params.
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
-
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
-
-        // Plan selection should succeed due to falling back on regular planning.
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
-
-        // Work the stage until it produces all results.
-        size_t numResults = 0;
-        PlanStage::StageState stageState = PlanStage::NEED_TIME;
-        while (stageState != PlanStage::IS_EOF) {
-            WorkingSetID id = WorkingSet::INVALID_ID;
-            stageState = subplan->work(&id);
-            ASSERT_NE(stageState, PlanStage::DEAD);
-            ASSERT_NE(stageState, PlanStage::FAILURE);
-
-            if (stageState == PlanStage::ADVANCED) {
-                ++numResults;
-                WorkingSetMember* member = ws.get(id);
-                ASSERT(member->hasObj());
-                ASSERT(SimpleBSONObjComparator::kInstance.evaluate(
-                           member->obj.value() == BSON("_id" << 1 << "a" << 1 << "b" << 2)) ||
-                       SimpleBSONObjComparator::kInstance.evaluate(
-                           member->obj.value() == BSON("_id" << 3 << "a" << 1 << "c" << 3)));
-            }
+TEST_F(QueryStageSubplanTest, QueryStageSubplanPlanContainedOr) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
+    addIndex(BSON("b" << 1 << "a" << 1));
+    addIndex(BSON("c" << 1 << "a" << 1));
+
+    BSONObj query = fromjson("{a: 1, $or: [{b: 2}, {c: 3}]}");
+
+    // Two of these documents match.
+    insert(BSON("_id" << 1 << "a" << 1 << "b" << 2));
+    insert(BSON("_id" << 2 << "a" << 2 << "b" << 2));
+    insert(BSON("_id" << 3 << "a" << 1 << "c" << 3));
+    insert(BSON("_id" << 4 << "a" << 1 << "c" << 4));
+
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(query);
+    auto cq = unittest::assertGet(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
+
+    Collection* collection = ctx.getCollection();
+
+    // Get planner params.
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
+
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
+
+    // Plan selection should succeed due to falling back on regular planning.
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+
+    // Work the stage until it produces all results.
+    size_t numResults = 0;
+    PlanStage::StageState stageState = PlanStage::NEED_TIME;
+    while (stageState != PlanStage::IS_EOF) {
+        WorkingSetID id = WorkingSet::INVALID_ID;
+        stageState = subplan->work(&id);
+        ASSERT_NE(stageState, PlanStage::DEAD);
+        ASSERT_NE(stageState, PlanStage::FAILURE);
+
+        if (stageState == PlanStage::ADVANCED) {
+            ++numResults;
+            WorkingSetMember* member = ws.get(id);
+            ASSERT(member->hasObj());
+            ASSERT(SimpleBSONObjComparator::kInstance.evaluate(
+                       member->obj.value() == BSON("_id" << 1 << "a" << 1 << "b" << 2)) ||
+                   SimpleBSONObjComparator::kInstance.evaluate(
+                       member->obj.value() == BSON("_id" << 3 << "a" << 1 << "c" << 3)));
         }
-
-        ASSERT_EQ(numResults, 2U);
     }
-};
+
+    ASSERT_EQ(numResults, 2U);
+}
 
 /**
  * Test the subplan stage's ability to answer a rooted $or query with a $ne and a sort.
  *
  * Regression test for SERVER-19388.
  */
-class QueryStageSubplanPlanRootedOrNE : public QueryStageSubplanBase {
-public:
-    void run() {
-        OldClientWriteContext ctx(&_opCtx, nss.ns());
-        addIndex(BSON("a" << 1 << "b" << 1));
-        addIndex(BSON("a" << 1 << "c" << 1));
-
-        // Every doc matches.
-        insert(BSON("_id" << 1 << "a" << 1));
-        insert(BSON("_id" << 2 << "a" << 2));
-        insert(BSON("_id" << 3 << "a" << 3));
-        insert(BSON("_id" << 4));
-
-        auto qr = stdx::make_unique<QueryRequest>(nss);
-        qr->setFilter(fromjson("{$or: [{a: 1}, {a: {$ne:1}}]}"));
-        qr->setSort(BSON("d" << 1));
-        auto cq = unittest::assertGet(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
-
-        Collection* collection = ctx.getCollection();
-
-        QueryPlannerParams plannerParams;
-        fillOutPlannerParams(&_opCtx, collection, cq.get(), &plannerParams);
-
-        WorkingSet ws;
-        std::unique_ptr<SubplanStage> subplan(
-            new SubplanStage(&_opCtx, collection, &ws, plannerParams, cq.get()));
-
-        PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
-        ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
-
-        size_t numResults = 0;
-        PlanStage::StageState stageState = PlanStage::NEED_TIME;
-        while (stageState != PlanStage::IS_EOF) {
-            WorkingSetID id = WorkingSet::INVALID_ID;
-            stageState = subplan->work(&id);
-            ASSERT_NE(stageState, PlanStage::DEAD);
-            ASSERT_NE(stageState, PlanStage::FAILURE);
-            if (stageState == PlanStage::ADVANCED) {
-                ++numResults;
-            }
+TEST_F(QueryStageSubplanTest, QueryStageSubplanPlanRootedOrNE) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
+    addIndex(BSON("a" << 1 << "b" << 1));
+    addIndex(BSON("a" << 1 << "c" << 1));
+
+    // Every doc matches.
+    insert(BSON("_id" << 1 << "a" << 1));
+    insert(BSON("_id" << 2 << "a" << 2));
+    insert(BSON("_id" << 3 << "a" << 3));
+    insert(BSON("_id" << 4));
+
+    auto qr = stdx::make_unique<QueryRequest>(nss);
+    qr->setFilter(fromjson("{$or: [{a: 1}, {a: {$ne:1}}]}"));
+    qr->setSort(BSON("d" << 1));
+    auto cq = unittest::assertGet(CanonicalQuery::canonicalize(opCtx(), std::move(qr)));
+
+    Collection* collection = ctx.getCollection();
+
+    QueryPlannerParams plannerParams;
+    fillOutPlannerParams(opCtx(), collection, cq.get(), &plannerParams);
+
+    WorkingSet ws;
+    std::unique_ptr<SubplanStage> subplan(
+        new SubplanStage(opCtx(), collection, &ws, plannerParams, cq.get()));
+
+    PlanYieldPolicy yieldPolicy(PlanExecutor::NO_YIELD, _clock);
+    ASSERT_OK(subplan->pickBestPlan(&yieldPolicy));
+
+    size_t numResults = 0;
+    PlanStage::StageState stageState = PlanStage::NEED_TIME;
+    while (stageState != PlanStage::IS_EOF) {
+        WorkingSetID id = WorkingSet::INVALID_ID;
+        stageState = subplan->work(&id);
+        ASSERT_NE(stageState, PlanStage::DEAD);
+        ASSERT_NE(stageState, PlanStage::FAILURE);
+        if (stageState == PlanStage::ADVANCED) {
+            ++numResults;
         }
-
-        ASSERT_EQ(numResults, 4U);
     }
-};
 
-class All : public Suite {
-public:
-    All() : Suite("query_stage_subplan") {}
-
-    void setupTests() {
-        add<QueryStageSubplanGeo2dOr>();
-        add<QueryStageSubplanPlanFromCache>();
-        add<QueryStageSubplanDontCacheZeroResults>();
-        add<QueryStageSubplanDontCacheTies>();
-        add<QueryStageSubplanCanUseSubplanning>();
-        add<QueryStageSubplanRewriteToRootedOr>();
-        add<QueryStageSubplanPlanContainedOr>();
-        add<QueryStageSubplanPlanRootedOrNE>();
+    ASSERT_EQ(numResults, 4U);
+}
+
+TEST_F(QueryStageSubplanTest, ShouldReportErrorIfExceedsTimeLimitDuringPlanning) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
+    // Build a query with a rooted $or.
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(BSON("$or" << BSON_ARRAY(BSON("p1" << 1) << BSON("p2" << 2))));
+    auto canonicalQuery =
+        uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest)));
+
+    // Add 4 indices: 2 for each predicate to choose from.
+    addIndex(BSON("p1" << 1 << "opt1" << 1));
+    addIndex(BSON("p1" << 1 << "opt2" << 1));
+    addIndex(BSON("p2" << 1 << "opt1" << 1));
+    addIndex(BSON("p2" << 1 << "opt2" << 1));
+    QueryPlannerParams params;
+    fillOutPlannerParams(opCtx(), ctx.getCollection(), canonicalQuery.get(), &params);
+
+    // Add some data so planning has to do some thinking.
+    for (int i = 0; i < 100; ++i) {
+        insert(BSON("_id" << i << "p1" << 1 << "p2" << 1));
+        insert(BSON("_id" << 2 * i << "p1" << 1 << "p2" << 2));
+        insert(BSON("_id" << 3 * i << "p1" << 2 << "p2" << 1));
+        insert(BSON("_id" << 4 * i << "p1" << 2 << "p2" << 2));
     }
-};
-
-SuiteInstance<All> all;
 
-}  // namespace QueryStageSubplan
+    // Create the SubplanStage.
+    WorkingSet workingSet;
+    SubplanStage subplanStage(
+        opCtx(), ctx.getCollection(), &workingSet, params, canonicalQuery.get());
+
+    AlwaysTimeOutYieldPolicy alwaysTimeOutPolicy(serviceContext()->getFastClockSource());
+    ASSERT_EQ(ErrorCodes::ExceededTimeLimit, subplanStage.pickBestPlan(&alwaysTimeOutPolicy));
+}
+
+TEST_F(QueryStageSubplanTest, ShouldReportErrorIfKilledDuringPlanning) {
+    OldClientWriteContext ctx(opCtx(), nss.ns());
+    // Build a query with a rooted $or.
+    auto queryRequest = stdx::make_unique<QueryRequest>(nss);
+    queryRequest->setFilter(BSON("$or" << BSON_ARRAY(BSON("p1" << 1) << BSON("p2" << 2))));
+    auto canonicalQuery =
+        uassertStatusOK(CanonicalQuery::canonicalize(opCtx(), std::move(queryRequest)));
+
+    // Add 4 indices: 2 for each predicate to choose from.
+    addIndex(BSON("p1" << 1 << "opt1" << 1));
+    addIndex(BSON("p1" << 1 << "opt2" << 1));
+    addIndex(BSON("p2" << 1 << "opt1" << 1));
+    addIndex(BSON("p2" << 1 << "opt2" << 1));
+    QueryPlannerParams params;
+    fillOutPlannerParams(opCtx(), ctx.getCollection(), canonicalQuery.get(), &params);
+
+    // Create the SubplanStage.
+    WorkingSet workingSet;
+    SubplanStage subplanStage(
+        opCtx(), ctx.getCollection(), &workingSet, params, canonicalQuery.get());
+
+    AlwaysPlanKilledYieldPolicy alwaysPlanKilledYieldPolicy(serviceContext()->getFastClockSource());
+    ASSERT_EQ(ErrorCodes::QueryPlanKilled, subplanStage.pickBestPlan(&alwaysPlanKilledYieldPolicy));
+}
+
+}  // namespace
+}  // namespace mongo