SERVER-21690 TextOrStage should fetch doc when RecordId first observed

3 files changed, 39 insertions, 29 deletions

diff --git a/src/mongo/db/exec/text.cpp b/src/mongo/db/exec/text.cpp
index 5c2f1fa1e67..424b372b017 100644
--- a/src/mongo/db/exec/text.cpp
+++ b/src/mongo/db/exec/text.cpp
@@ -30,7 +30,6 @@
 
 #include <vector>
 
-#include "mongo/db/exec/fetch.h"
 #include "mongo/db/exec/filter.h"
 #include "mongo/db/exec/index_scan.h"
 #include "mongo/db/exec/text_or.h"
@@ -135,11 +134,8 @@ unique_ptr<PlanStage> TextStage::buildTextTree(OperationContext* txn,
         textScorer->addChild(make_unique<IndexScan>(txn, ixparams, ws, nullptr));
     }
 
-    auto fetcher = make_unique<FetchStage>(
-        txn, ws, textScorer.release(), nullptr, _params.index->getCollection());
-
     auto matcher =
-        make_unique<TextMatchStage>(txn, std::move(fetcher), _params.query, _params.spec, ws);
+        make_unique<TextMatchStage>(txn, std::move(textScorer), _params.query, _params.spec, ws);
 
     unique_ptr<PlanStage> treeRoot = std::move(matcher);
     return treeRoot;
diff --git a/src/mongo/db/exec/text_or.cpp b/src/mongo/db/exec/text_or.cpp
index a022a7cce0d..66ac3c51953 100644
--- a/src/mongo/db/exec/text_or.cpp
+++ b/src/mongo/db/exec/text_or.cpp
@@ -351,18 +351,21 @@ PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out)
     invariant(wsm->getState() == WorkingSetMember::LOC_AND_IDX);
     invariant(1 == wsm->keyData.size());
     const IndexKeyDatum newKeyData = wsm->keyData.back();  // copy to keep it around.
-
     TextRecordData* textRecordData = &_scores[wsm->loc];
-    double* documentAggregateScore = &textRecordData->score;
 
-    if (WorkingSet::INVALID_ID == textRecordData->wsid) {
-        // We haven't seen this RecordId before. Keep the working set member around (it may be
-        // force-fetched on saveState()).
-        textRecordData->wsid = wsid;
+    if (textRecordData->score < 0) {
+        // We have already rejected this document for not matching the filter.
+        invariant(WorkingSet::INVALID_ID == textRecordData->wsid);
+        _ws->free(wsid);
+        return NEED_TIME;
+    }
 
+    if (WorkingSet::INVALID_ID == textRecordData->wsid) {
+        // We haven't seen this RecordId before.
+        invariant(textRecordData->score == 0);
+        bool shouldKeep = true;
         if (_filter) {
             // We have not seen this document before and need to apply a filter.
-            bool shouldKeep;
             bool wasDeleted = false;
             try {
                 TextMatchableDocument tdoc(getOpCtx(),
@@ -386,18 +389,35 @@ PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out)
                 wasDeleted = true;
             }
 
-            if (!shouldKeep) {
-                if (wasDeleted || wsm->hasObj()) {
-                    // We had to fetch but we're not going to return it.
-                    ++_specificStats.fetches;
-                }
-                _ws->free(textRecordData->wsid);
-                textRecordData->wsid = WorkingSet::INVALID_ID;
-                *documentAggregateScore = -1;
-                return NEED_TIME;
+            if (wasDeleted || wsm->hasObj()) {
+                ++_specificStats.fetches;
+            }
+        }
+
+        if (shouldKeep && !wsm->hasObj()) {
+            // Our parent expects LOC_AND_OBJ members, so we fetch the document here if we haven't
+            // already.
+            try {
+                shouldKeep = WorkingSetCommon::fetch(getOpCtx(), _ws, wsid, _recordCursor);
+                ++_specificStats.fetches;
+            } catch (const WriteConflictException& wce) {
+                wsm->makeObjOwnedIfNeeded();
+                _idRetrying = wsid;
+                *out = WorkingSet::INVALID_ID;
+                return NEED_YIELD;
             }
         }
 
+        if (!shouldKeep) {
+            _ws->free(wsid);
+            textRecordData->score = -1;
+            return NEED_TIME;
+        }
+
+        textRecordData->wsid = wsid;
+
+        // Ensure that the BSONObj underlying the WorkingSetMember is owned in case we yield.
+        wsm->makeObjOwnedIfNeeded();
     } else {
         // We already have a working set member for this RecordId. Free the new WSM and retrieve the
         // old one. Note that since we don't keep all index keys, we could get a score that doesn't
@@ -408,11 +428,6 @@ PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out)
         wsm = _ws->get(textRecordData->wsid);
     }
 
-    if (*documentAggregateScore < 0) {
-        // We have already rejected this document for not matching the filter.
-        return NEED_TIME;
-    }
-
     // Locate score within possibly compound key: {prefix,term,score,suffix}.
     BSONObjIterator keyIt(newKeyData.keyData);
     for (unsigned i = 0; i < _ftsSpec.numExtraBefore(); i++) {
@@ -425,7 +440,7 @@ PlanStage::StageState TextOrStage::addTerm(WorkingSetID wsid, WorkingSetID* out)
     double documentTermScore = scoreElement.number();
 
     // Aggregate relevance score, term keys.
-    *documentAggregateScore += documentTermScore;
+    textRecordData->score += documentTermScore;
     return NEED_TIME;
 }
 
diff --git a/src/mongo/db/exec/text_or.h b/src/mongo/db/exec/text_or.h
index 68b98f0cb51..ade42440cc1 100644
--- a/src/mongo/db/exec/text_or.h
+++ b/src/mongo/db/exec/text_or.h
@@ -51,8 +51,7 @@ class OperationContext;
  * A blocking stage that returns the set of WSMs with RecordIDs of all of the documents that contain
  * the positive terms in the search query, as well as their scores.
  *
- * The WorkingSetMembers returned are in the LOC_AND_IDX state. If a filter is passed in, some
- * WorkingSetMembers may be returned in the LOC_AND_OBJ state.
+ * The WorkingSetMembers returned are fetched and in the LOC_AND_OBJ state.
  */
 class TextOrStage final : public PlanStage {
 public: