SERVER-16675 force fetch RecordIds buffered by the query system on saveState()

This fixes an issue with WiredTiger query isolation.

1 files changed, 44 insertions, 20 deletions

diff --git a/src/mongo/db/exec/text.cpp b/src/mongo/db/exec/text.cpp
index 037b20c28aa..95053c8c8da 100644
--- a/src/mongo/db/exec/text.cpp
+++ b/src/mongo/db/exec/text.cpp
@@ -224,8 +224,7 @@ namespace mongo {
             invariant(1 == wsm->keyData.size());
             invariant(wsm->hasLoc());
             IndexKeyDatum& keyDatum = wsm->keyData.back();
-            addTerm(keyDatum.keyData, wsm->loc);
-            _ws->free(id);
+            addTerm(keyDatum.keyData, id);
             return PlanStage::NEED_TIME;
         }
         else if (PlanStage::IS_EOF == childState) {
@@ -270,31 +269,40 @@ namespace mongo {
         }
 
         // Filter for phrases and negative terms, score and truncate.
-        RecordId loc = _scoreIterator->first;
-        double score = _scoreIterator->second;
+        TextRecordData textRecordData = _scoreIterator->second;
+        WorkingSetMember* wsm = _ws->get(textRecordData.wsid);
         _scoreIterator++;
 
         // Ignore non-matched documents.
-        if (score < 0) {
+        if (textRecordData.score < 0) {
+            _ws->free(textRecordData.wsid);
             return PlanStage::NEED_TIME;
         }
 
-        // Fetch the document
-        BSONObj doc(_params.index->getCollection()->docFor(_txn, loc));
+        // Retrieve the document. We may already have the document due to force-fetching before
+        // a yield. If not, then we fetch the document here.
+        BSONObj doc;
+        if (wsm->hasObj()) {
+            doc = wsm->obj;
+        }
+        else {
+            doc = _params.index->getCollection()->docFor(_txn, wsm->loc);
+            wsm->obj = doc;
+            wsm->keyData.clear();
+            wsm->state = WorkingSetMember::LOC_AND_UNOWNED_OBJ;
+        }
 
         // Filter for phrases and negated terms
         if (_params.query.hasNonTermPieces()) {
             if (!_ftsMatcher.matchesNonTerm(doc)) {
+                _ws->free(textRecordData.wsid);
                 return PlanStage::NEED_TIME;
             }
         }
 
-        *out = _ws->allocate();
-        WorkingSetMember* member = _ws->get(*out);
-        member->loc = loc;
-        member->obj = doc;
-        member->state = WorkingSetMember::LOC_AND_UNOWNED_OBJ;
-        member->addComputed(new TextScoreComputedData(score));
+        // Populate the working set member with the text score and return it.
+        wsm->addComputed(new TextScoreComputedData(textRecordData.score));
+        *out = textRecordData.wsid;
         return PlanStage::ADVANCED;
     }
 
@@ -356,8 +364,24 @@ namespace mongo {
         bool* _fetched;
     };
 
-    void TextStage::addTerm(const BSONObj& key, const RecordId& loc) {
-        double *documentAggregateScore = &_scores[loc];
+    void TextStage::addTerm(const BSONObj key, WorkingSetID wsid) {
+        WorkingSetMember* wsm = _ws->get(wsid);
+        TextRecordData* textRecordData = &_scores[wsm->loc];
+
+        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;
+        }
+        else {
+            // We already have a working set member for this RecordId. Free the old
+            // WSM and retrieve the new one.
+            invariant(wsid != textRecordData->wsid);
+            _ws->free(wsid);
+            wsm = _ws->get(textRecordData->wsid);
+        }
+
+        double* documentAggregateScore = &textRecordData->score;
 
         ++_specificStats.keysExamined;
 
@@ -371,7 +395,7 @@ namespace mongo {
 
         BSONElement scoreElement = keyIt.next();
         double documentTermScore = scoreElement.number();
-        
+
         // Handle filtering.
         if (*documentAggregateScore < 0) {
             // We have already rejected this document.
@@ -383,10 +407,10 @@ namespace mongo {
                 // We have not seen this document before and need to apply a filter.
                 bool fetched = false;
                 TextMatchableDocument tdoc(_txn,
-                                           _params.index->keyPattern(), 
-                                           key, 
-                                           loc, 
-                                           _params.index->getCollection(), 
+                                           _params.index->keyPattern(),
+                                           key,
+                                           wsm->loc,
+                                           _params.index->getCollection(),
                                            &fetched);
 
                 if (!_filter->matches(&tdoc)) {