SERVER-17198 prevent invalid logOp rollback in findAndModify

(cherry picked from commit 570951951f752d1a31bb9e7f0d4276c1b533f25e)

8 files changed, 68 insertions, 30 deletions

diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index d12c2df3704..68b0bf804d2 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -393,6 +393,7 @@ namespace mongo {
                     request.setUpdates(update);
                     request.setUpsert(upsert);
                     request.setUpdateOpLog();
+                    request.setStoreResultDoc(returnNew);
 
                     request.setYieldPolicy(PlanExecutor::YIELD_MANUAL);
 
@@ -420,30 +421,11 @@ namespace mongo {
                     }
 
                     LOG(3) << "update result: "  << res ;
-                    if ( returnNew ) {
-                        if ( !res.upserted.isEmpty() ) {
-                            BSONElement upsertedElem = res.upserted[kUpsertedFieldName];
-                            LOG(3) << "using new _id to get new doc: "
-                                   << upsertedElem;
-                            queryModified = upsertedElem.wrap("_id");
-                        }
-                        else if ( queryModified["_id"].type() ) {
-                            // we do this so that if the update changes the fields, it still matches
-                            queryModified = queryModified["_id"].wrap();
-                        }
-
-                        LOG(3) << "using modified query to return the new doc: " << queryModified;
-                        if ( ! Helpers::findOne( txn, collection, queryModified, doc ) ) {
-                            errmsg = str::stream() << "can't find object after modification  " 
-                                                   << " ns: " << ns 
-                                                   << " queryModified: " << queryModified 
-                                                   << " queryOriginal: " << query;
-                            log() << errmsg << endl;
-                            return false;
-                        }
-                        _appendHelper(result, doc, true, fields, whereCallback);
+                    if (returnNew) {
+                        dassert(!res.newObj.isEmpty());
+                        _appendHelper(result, res.newObj, true, fields, whereCallback);
                     }
-                    
+
                     BSONObjBuilder le( result.subobjStart( "lastErrorObject" ) );
                     le.appendBool( "updatedExisting" , res.existing );
                     le.appendNumber( "n" , res.numMatched );
@@ -451,7 +433,6 @@ namespace mongo {
                         le.append( res.upserted[kUpsertedFieldName] );
                     }
                     le.done();
-                    
                 }
             }
 
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index 59ce77035cb..252122f0025 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -632,6 +632,13 @@ namespace mongo {
         // The object that was inserted. This is an empty document if no insert was performed.
         BSONObj objInserted;
 
+        // The document resulting from the update (or the document inserted for an upsert). This is
+        // an empty document if no existing doc was modified and no insert happened, or if this is
+        // a multi-update.
+        //
+        // Only set if requested via UpdateRequest::setStoreResultDoc().
+        BSONObj newObj;
+
         // Invalidated documents can be force-fetched, causing the now invalid RecordId to
         // be thrown out. The update stage skips over any results which do not have the
         // RecordId to update.
diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index c79925cb8c9..329901bfe77 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -519,6 +519,11 @@ namespace mongo {
             docWasModified = false;
         }
 
+        if (!docWasModified && request->shouldStoreResultDoc()) {
+            // The update is a no-op, so the result doc is the same as the old doc.
+            _specificStats.newObj = oldObj.value().getOwned();
+        }
+
         if (docWasModified) {
 
             // Verify that no immutable fields were changed and data is valid for storage.
@@ -602,6 +607,11 @@ namespace mongo {
 
             invariant(oldObj.snapshotId() == _txn->recoveryUnit()->getSnapshotId());
             wunit.commit();
+
+            if (request->shouldStoreResultDoc()) {
+                // We just committed a single update. Hold onto the resulting document.
+                _specificStats.newObj = newObj.getOwned();
+            }
         }
 
         // Only record doc modifications if they wrote (exclude no-ops). Explains get
@@ -683,6 +693,9 @@ namespace mongo {
                 newObj.objsize() <= BSONObjMaxUserSize);
 
         _specificStats.objInserted = newObj;
+        if (request->shouldStoreResultDoc()) {
+            _specificStats.newObj = newObj;
+        }
 
         // If this is an explain, bail out now without doing the insert.
         if (request->isExplain()) {
@@ -1009,7 +1022,8 @@ namespace mongo {
                             !updateStats->isDocReplacement /* $mod or obj replacement */,
                             opDebug->nModified /* number of modified docs, no no-ops */,
                             opDebug->nMatched /* # of docs matched/updated, even no-ops */,
-                            updateStats->objInserted);
+                            updateStats->objInserted,
+                            updateStats->newObj);
     };
 
 } // namespace mongo
diff --git a/src/mongo/db/ops/parsed_update.cpp b/src/mongo/db/ops/parsed_update.cpp
index 948aa775084..8570ee212d8 100644
--- a/src/mongo/db/ops/parsed_update.cpp
+++ b/src/mongo/db/ops/parsed_update.cpp
@@ -42,6 +42,10 @@ namespace mongo {
         _canonicalQuery() { }
 
     Status ParsedUpdate::parseRequest() {
+        // It is invalid to request that the update plan stores a copy of the resulting document
+        // if it is a multi-update.
+        invariant(!(_request->shouldStoreResultDoc() && _request->isMulti()));
+
         // We parse the update portion before the query portion because the dispostion of the update
         // may determine whether or not we need to produce a CanonicalQuery at all.  For example, if
         // the update involves the positional-dollar operator, we must have a CanonicalQuery even if
diff --git a/src/mongo/db/ops/update_request.h b/src/mongo/db/ops/update_request.h
index 5c2692ca3ed..99649663ece 100644
--- a/src/mongo/db/ops/update_request.h
+++ b/src/mongo/db/ops/update_request.h
@@ -53,6 +53,7 @@ namespace mongo {
             , _fromReplication(false)
             , _lifecycle(NULL)
             , _isExplain(false)
+            , _storeResultDoc(false)
             , _yieldPolicy(PlanExecutor::YIELD_MANUAL) {}
 
         const NamespaceString& getNamespaceString() const {
@@ -142,6 +143,14 @@ namespace mongo {
             return _isExplain;
         }
 
+        inline void setStoreResultDoc(bool value = true) {
+            _storeResultDoc = value;
+        }
+
+        inline bool shouldStoreResultDoc() const {
+            return _storeResultDoc;
+        }
+
         inline void setYieldPolicy(PlanExecutor::YieldPolicy yieldPolicy) {
             _yieldPolicy = yieldPolicy;
         }
@@ -199,6 +208,14 @@ namespace mongo {
         // Whether or not we are requesting an explained update. Explained updates are read-only.
         bool _isExplain;
 
+        // Whether or not we keep an owned copy of the resulting document for a non-multi update.
+        // This allows someone executing an update to retrieve the resulting document without
+        // another query once the update is complete.
+        //
+        // It is illegal to use this flag in combination with the '_multi' flag, and doing so will
+        // trigger an invariant check.
+        bool _storeResultDoc;
+
         // Whether or not the update should yield. Defaults to YIELD_MANUAL.
         PlanExecutor::YieldPolicy _yieldPolicy;
 
diff --git a/src/mongo/db/ops/update_result.cpp b/src/mongo/db/ops/update_result.cpp
index f5e31c1a3d8..d2664d34e38 100644
--- a/src/mongo/db/ops/update_result.cpp
+++ b/src/mongo/db/ops/update_result.cpp
@@ -42,11 +42,13 @@ namespace mongo {
                                bool modifiers_,
                                unsigned long long numDocsModified_,
                                unsigned long long numMatched_,
-                               const BSONObj& upsertedObject_)
+                               const BSONObj& upsertedObject_,
+                               const BSONObj& newObj_)
         : existing(existing_),
           modifiers(modifiers_),
           numDocsModified(numDocsModified_),
-          numMatched(numMatched_) {
+          numMatched(numMatched_),
+          newObj(newObj_) {
 
         BSONElement id = upsertedObject_["_id"];
         if ( ! existing && numMatched == 1 && !id.eoo() ) {
diff --git a/src/mongo/db/ops/update_result.h b/src/mongo/db/ops/update_result.h
index 276813ae150..250b9d1fe23 100644
--- a/src/mongo/db/ops/update_result.h
+++ b/src/mongo/db/ops/update_result.h
@@ -43,7 +43,8 @@ namespace mongo {
                       bool modifiers_,
                       unsigned long long numDocsModified_,
                       unsigned long long numMatched_,
-                      const BSONObj& upsertedObject_ );
+                      const BSONObj& upsertedObject_,
+                      const BSONObj& newObj_ );
 
 
         // if existing objects were modified
@@ -61,13 +62,18 @@ namespace mongo {
         // if something was upserted, the new _id of the object
         BSONObj upserted;
 
+        // For a non-multi update, the new version of the document. If we did an insert, this
+        // is the full document that got inserted (whereas 'upserted' is just the _id field).
+        BSONObj newObj;
+
         const std::string toString() const {
             return str::stream()
                         << " upserted: " << upserted
                         << " modifiers: " << modifiers
                         << " existing: " << existing
                         << " numDocsModified: " << numDocsModified
-                        << " numMatched: " << numMatched;
+                        << " numMatched: " << numMatched
+                        << " newObj: " << newObj;
         }
     };
 
diff --git a/src/mongo/db/storage/in_memory/in_memory_record_store.cpp b/src/mongo/db/storage/in_memory/in_memory_record_store.cpp
index e4fb801032a..929a4f51691 100644
--- a/src/mongo/db/storage/in_memory/in_memory_record_store.cpp
+++ b/src/mongo/db/storage/in_memory/in_memory_record_store.cpp
@@ -329,7 +329,14 @@ namespace mongo {
     }
 
     bool InMemoryRecordStore::updateWithDamagesSupported() const {
-        return true;
+        // TODO: Currently the UpdateStage assumes that updateWithDamages will apply the
+        // damages directly to the unowned BSONObj containing the record to be modified.
+        // The implementation of updateWithDamages() below copies the old record to a
+        // a new one and then applies the damages.
+        //
+        // We should be able to enable updateWithDamages() here once this assumption is
+        // relaxed.
+        return false;
     }
 
     Status InMemoryRecordStore::updateWithDamages( OperationContext* txn,