SERVER-17062 Make query execution handle WriteConflictExceptions where possible

1 files changed, 40 insertions, 33 deletions

diff --git a/src/mongo/db/exec/fetch.cpp b/src/mongo/db/exec/fetch.cpp
index 9d78111ea35..307804dd58f 100644
--- a/src/mongo/db/exec/fetch.cpp
+++ b/src/mongo/db/exec/fetch.cpp
@@ -26,9 +26,12 @@
  *    it in the license file.
  */
 
+#include "mongo/platform/basic.h"
+
 #include "mongo/db/exec/fetch.h"
 
 #include "mongo/db/catalog/collection.h"
+#include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/exec/filter.h"
 #include "mongo/db/exec/scoped_timer.h"
 #include "mongo/db/exec/working_set_common.h"
@@ -54,13 +57,13 @@ namespace mongo {
           _ws(ws),
           _child(child),
           _filter(filter),
-          _idBeingPagedIn(WorkingSet::INVALID_ID),
+          _idRetrying(WorkingSet::INVALID_ID),
           _commonStats(kStageType) { }
 
     FetchStage::~FetchStage() { }
 
     bool FetchStage::isEOF() {
-        if (WorkingSet::INVALID_ID != _idBeingPagedIn) {
+        if (WorkingSet::INVALID_ID != _idRetrying) {
             // We asked the parent for a page-in, but still haven't had a chance to return the
             // paged in document
             return false;
@@ -77,21 +80,17 @@ namespace mongo {
 
         if (isEOF()) { return PlanStage::IS_EOF; }
 
-        // We might have a fetched result to return.
-        if (WorkingSet::INVALID_ID != _idBeingPagedIn) {
-            WorkingSetID id = _idBeingPagedIn;
-            _idBeingPagedIn = WorkingSet::INVALID_ID;
-            WorkingSetMember* member = _ws->get(id);
-
-            WorkingSetCommon::completeFetch(_txn, member, _collection);
-
-            return returnIfMatches(member, id, out);
+        // Either retry the last WSM we worked on or get a new one from our child.
+        WorkingSetID id;
+        StageState status;
+        if (_idRetrying == WorkingSet::INVALID_ID) {
+            status = _child->work(&id);
+        }
+        else {
+            status = ADVANCED;
+            id = _idRetrying;
+            _idRetrying = WorkingSet::INVALID_ID;
         }
-
-        // If we're here, we're not waiting for a RecordId to be fetched.  Get another to-be-fetched
-        // result from our child.
-        WorkingSetID id = WorkingSet::INVALID_ID;
-        StageState status = _child->work(&id);
 
         if (PlanStage::ADVANCED == status) {
             WorkingSetMember* member = _ws->get(id);
@@ -107,25 +106,33 @@ namespace mongo {
 
                 // We might need to retrieve 'nextLoc' from secondary storage, in which case we send
                 // a NEED_FETCH request up to the PlanExecutor.
-                if (!member->loc.isNull()) {
-                    std::auto_ptr<RecordFetcher> fetcher(
-                        _collection->documentNeedsFetch(_txn, member->loc));
-                    if (NULL != fetcher.get()) {
-                        // There's something to fetch. Hand the fetcher off to the WSM, and pass up
-                        // a fetch request.
-                        _idBeingPagedIn = id;
-                        member->setFetcher(fetcher.release());
-                        *out = id;
-                        _commonStats.needFetch++;
-                        return NEED_FETCH;
-                    }
+                std::auto_ptr<RecordFetcher> fetcher(_collection->documentNeedsFetch(_txn,
+                                                                                     member->loc));
+                if (NULL != fetcher.get()) {
+                    // There's something to fetch. Hand the fetcher off to the WSM, and pass up
+                    // a fetch request.
+                    _idRetrying = id;
+                    member->setFetcher(fetcher.release());
+                    *out = id;
+                    _commonStats.needFetch++;
+                    return NEED_FETCH;
                 }
 
                 // The doc is already in memory, so go ahead and grab it. Now we have a RecordId
                 // as well as an unowned object
-                member->obj = _collection->docFor(_txn, member->loc);
-                member->keyData.clear();
-                member->state = WorkingSetMember::LOC_AND_UNOWNED_OBJ;
+                try {
+                    if (!WorkingSetCommon::fetch(_txn, member, _collection)) {
+                        _ws->free(id);
+                        _commonStats.needTime++;
+                        return NEED_TIME;
+                    }
+                }
+                catch (const WriteConflictException& wce) {
+                    _idRetrying = id;
+                    *out = WorkingSet::INVALID_ID;
+                    _commonStats.needFetch++;
+                    return NEED_FETCH;
+                }
             }
 
             return returnIfMatches(member, id, out);
@@ -174,8 +181,8 @@ namespace mongo {
 
         // It's possible that the loc getting invalidated is the one we're about to
         // fetch. In this case we do a "forced fetch" and put the WSM in owned object state.
-        if (WorkingSet::INVALID_ID != _idBeingPagedIn) {
-            WorkingSetMember* member = _ws->get(_idBeingPagedIn);
+        if (WorkingSet::INVALID_ID != _idRetrying) {
+            WorkingSetMember* member = _ws->get(_idRetrying);
             if (member->hasLoc() && (member->loc == dl)) {
                 // Fetch it now and kill the diskloc.
                 WorkingSetCommon::fetchAndInvalidateLoc(txn, member, _collection);