SERVER-14425: handle deadlocks in the middle up a multi-update

6 files changed, 63 insertions, 17 deletions

diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index e3ece80865a..730bbe2b48d 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -33,6 +33,7 @@
 #include "mongo/db/exec/update.h"
 
 #include "mongo/bson/mutable/algorithm.h"
+#include "mongo/db/concurrency/deadlock.h"
 #include "mongo/db/exec/scoped_timer.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/ops/update_lifecycle.h"
@@ -488,10 +489,6 @@ namespace mongo {
             }
         }
 
-
-        // Save state before making changes
-        saveState();
-
         {
             WriteUnitOfWork wunit(request->getOpCtx());
 
@@ -559,14 +556,6 @@ namespace mongo {
             wunit.commit();
         }
 
-
-        // Restore state after modification
-
-        // As restoreState may restore (recreate) cursors, make sure to restore the
-        // state outside of the WritUnitOfWork.
-
-        restoreState(request->getOpCtx());
-
         // Only record doc modifications if they wrote (exclude no-ops). Explains get
         // recorded as if they wrote.
         if (docWasModified) {
@@ -753,8 +742,45 @@ namespace mongo {
 
             ++_specificStats.nMatched;
 
+            // Save state before making changes
+            saveState();
+
             // Do the update and return.
-            transformAndUpdate(oldObj, loc);
+            BSONObj reFetched;
+            while ( 1 ) {
+                try {
+                    transformAndUpdate(reFetched.isEmpty() ? oldObj : reFetched , loc);
+                    break;
+                }
+                catch ( const DeadLockException& de ) {
+                    if ( !_params.request->isMulti() ) {
+                        // for single cases, we just restart.
+                        throw;
+                    }
+
+                    log() << "got deadlock in the middle of a multi-update, redoing the doc";
+                    OperationContext* txn = _params.request->getOpCtx();
+                    txn->recoveryUnit()->commitAndRestart();
+                    if ( !_collection->findDoc( txn, loc, &reFetched ) ) {
+                        // document was deleted, we're done here
+                        break;
+                    }
+                    // we have to re-match the doc as it might not match anymore
+                    if ( !_params.canonicalQuery->root()->matchesBSON( reFetched, NULL ) ) {
+                        // doesn't match!
+                        break;
+                    }
+                    // now we try again!
+                }
+            }
+
+            // Restore state after modification
+
+            // As restoreState may restore (recreate) cursors, make sure to restore the
+            // state outside of the WritUnitOfWork.
+
+            restoreState(_params.request->getOpCtx());
+
             ++_commonStats.needTime;
             return PlanStage::NEED_TIME;
         }
diff --git a/src/mongo/db/ops/update_executor.cpp b/src/mongo/db/ops/update_executor.cpp
index 2975b67f745..d403ea9edeb 100644
--- a/src/mongo/db/ops/update_executor.cpp
+++ b/src/mongo/db/ops/update_executor.cpp
@@ -270,6 +270,7 @@ namespace mongo {
                                                      &cqRaw,
                                                      whereCallback);
         if (status.isOK()) {
+            cqRaw->setIsForWrite( true );
             _canonicalQuery.reset(cqRaw);
             _isQueryParsed = true;
         }
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index cca9dd79e3b..92650f7f30e 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -531,6 +531,7 @@ namespace mongo {
     Status CanonicalQuery::init(LiteParsedQuery* lpq,
                                 const MatchExpressionParser::WhereCallback& whereCallback,
                                 MatchExpression* root) {
+        _isForWrite = false;
         _pq.reset(lpq);
 
         // Normalize, sort and validate tree.
diff --git a/src/mongo/db/query/canonical_query.h b/src/mongo/db/query/canonical_query.h
index ff66f6a18fb..800f2c833a0 100644
--- a/src/mongo/db/query/canonical_query.h
+++ b/src/mongo/db/query/canonical_query.h
@@ -172,6 +172,9 @@ namespace mongo {
         std::string toString() const;
         std::string toStringShort() const;
 
+        bool isForWrite() const { return _isForWrite; }
+        void setIsForWrite( bool w ) { _isForWrite = w; }
+
         /**
          * Validates match expression, checking for certain
          * combinations of operators in match expression and
@@ -210,7 +213,6 @@ namespace mongo {
          * while exploring the enumeration space we do it here.
          */
         static MatchExpression* logicalRewrite(MatchExpression* tree);
-
     private:
         // You must go through canonicalize to create a CanonicalQuery.
         CanonicalQuery() { }
@@ -240,6 +242,8 @@ namespace mongo {
          * for minimal user comprehension.
          */
         PlanCacheKey _cacheKey;
+
+        bool _isForWrite;
     };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp
index 13d978f8d3d..4c6c64f1bd0 100644
--- a/src/mongo/db/query/planner_access.cpp
+++ b/src/mongo/db/query/planner_access.cpp
@@ -450,7 +450,10 @@ namespace mongo {
 
     // static
     bool QueryPlannerAccess::orNeedsFetch(const ScanBuildingState* scanState) {
-        if (scanState->loosestBounds == IndexBoundsBuilder::EXACT) {
+        if (!scanState->canTrimExpression) {
+            return true;
+        }
+        else if (scanState->loosestBounds == IndexBoundsBuilder::EXACT) {
             return false;
         }
         else if (scanState->loosestBounds == IndexBoundsBuilder::INEXACT_FETCH) {
@@ -599,6 +602,10 @@ namespace mongo {
                                                std::vector<QuerySolutionNode*>* out) {
         // Initialize the ScanBuildingState.
         ScanBuildingState scanState(root, inArrayOperator, indices);
+        if ( query.isForWrite() ) {
+            // A write may need to re-apply if it has to re-fetch a document after a deadlock.
+            scanState.canTrimExpression = false;
+        }
 
         while (scanState.curChild < root->numChildren()) {
             MatchExpression* child = root->getChild(scanState.curChild);
@@ -1082,7 +1089,7 @@ namespace mongo {
                 // superset of documents that satisfy the predicate, and we must check the
                 // predicate.
 
-                if (tightness == IndexBoundsBuilder::EXACT) {
+                if (tightness == IndexBoundsBuilder::EXACT && !query.isForWrite()) {
                     return soln;
                 }
                 else if (tightness == IndexBoundsBuilder::INEXACT_COVERED
@@ -1273,7 +1280,10 @@ namespace mongo {
         MatchExpression* child = root->getChild(scanState->curChild);
         const IndexEntry& index = scanState->indices[scanState->currentIndexNumber];
 
-        if (scanState->inArrayOperator) {
+        if (!scanState->canTrimExpression) {
+            ++scanState->curChild;
+        }
+        else if (scanState->inArrayOperator) {
             // We're inside an array operator. The entire array operator expression
             // should always be affixed as a filter. We keep 'curChild' in the $and
             // for affixing later.
diff --git a/src/mongo/db/query/planner_access.h b/src/mongo/db/query/planner_access.h
index b512a25880c..9ff57d6b941 100644
--- a/src/mongo/db/query/planner_access.h
+++ b/src/mongo/db/query/planner_access.h
@@ -106,6 +106,7 @@ namespace mongo {
                               const std::vector<IndexEntry>& indexList)
                 : root(theRoot),
                   inArrayOperator(inArrayOp),
+                  canTrimExpression(true),
                   indices(indexList),
                   currentScan(NULL),
                   curChild(0),
@@ -139,6 +140,9 @@ namespace mongo {
             // Are we inside an array operator such as $elemMatch or $all?
             bool inArrayOperator;
 
+            // Are we allowed to trip the match expression for indexes, or should we re-eval.
+            bool canTrimExpression;
+
             // A list of relevant indices which 'root' may be tagged to use.
             const std::vector<IndexEntry>& indices;