SERVER-17303 findAndModify upsert calls Collection::insertDocument() directly

2 files changed, 97 insertions, 42 deletions

diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index a2b1c991185..04d7813cb47 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -424,8 +424,8 @@ namespace mongo {
             }
 
             return Status::OK();
-
         }
+
     } // namespace
 
     // static
@@ -627,14 +627,15 @@ namespace mongo {
         }
     }
 
-    void UpdateStage::doInsert() {
-        _specificStats.inserted = true;
-
-        const UpdateRequest* request = _params.request;
-        UpdateDriver* driver = _params.driver;
-        CanonicalQuery* cq = _params.canonicalQuery;
-        UpdateLifecycle* lifecycle = request->getLifecycle();
-
+    // static
+    Status UpdateStage::applyUpdateOpsForInsert(const CanonicalQuery* cq,
+                                                const BSONObj& query,
+                                                UpdateDriver* driver,
+                                                UpdateLifecycle* lifecycle,
+                                                mutablebson::Document* doc,
+                                                bool isInternalRequest,
+                                                UpdateStats* stats,
+                                                BSONObj* out) {
         // Since this is an insert (no docs found and upsert:true), we will be logging it
         // as an insert in the oplog. We don't need the driver's help to build the
         // oplog record, then. We also set the context of the update driver to the INSERT_CONTEXT.
@@ -642,61 +643,88 @@ namespace mongo {
         driver->setLogOp(false);
         driver->setContext(ModifierInterface::ExecInfo::INSERT_CONTEXT);
 
-        // Reset the document we will be writing to
-        _doc.reset();
-
-        // The original document we compare changes to - immutable paths must not change
-        BSONObj original;
-
-        bool isInternalRequest = request->isFromReplication() || request->isFromMigration();
-
         const vector<FieldRef*>* immutablePaths = NULL;
         if (!isInternalRequest && lifecycle)
             immutablePaths = lifecycle->getImmutableFields();
 
-        // Calling populateDocumentWithQueryFields will populate the '_doc' with fields from the
-        // query which creates the base of the update for the inserted doc (because upsert
-        // was true).
+        // The original document we compare changes to - immutable paths must not change
+        BSONObj original;
+
         if (cq) {
-            uassertStatusOK(driver->populateDocumentWithQueryFields(cq, immutablePaths, _doc));
+            Status status = driver->populateDocumentWithQueryFields(cq, immutablePaths, *doc);
+            if (!status.isOK()) {
+                return status;
+            }
+
             if (driver->isDocReplacement())
-                _specificStats.fastmodinsert = true;
-            original = _doc.getObject();
+                stats->fastmodinsert = true;
+            original = doc->getObject();
         }
         else {
-            fassert(17354, CanonicalQuery::isSimpleIdQuery(request->getQuery()));
-            BSONElement idElt = request->getQuery()[idFieldName];
+            fassert(17354, CanonicalQuery::isSimpleIdQuery(query));
+            BSONElement idElt = query[idFieldName];
             original = idElt.wrap();
-            fassert(17352, _doc.root().appendElement(idElt));
+            fassert(17352, doc->root().appendElement(idElt));
         }
 
-        // Apply the update modifications and then log the update as an insert manually.
-        Status status = driver->update(StringData(), &_doc);
-        if (!status.isOK()) {
-            uasserted(16836, status.reason());
+        // Apply the update modifications here.
+        Status updateStatus = driver->update(StringData(), doc);
+        if (!updateStatus.isOK()) {
+            return Status(updateStatus.code(), updateStatus.reason(), 16836);
         }
 
         // Ensure _id exists and is first
-        uassertStatusOK(ensureIdAndFirst(_doc));
+        Status idAndFirstStatus = ensureIdAndFirst(*doc);
+        if (!idAndFirstStatus.isOK()) {
+            return idAndFirstStatus;
+        }
 
         // Validate that the object replacement or modifiers resulted in a document
         // that contains all the immutable keys and can be stored if it isn't coming
         // from a migration or via replication.
-        if (!isInternalRequest){
+        if (!isInternalRequest) {
             FieldRefSet noFields;
             // This will only validate the modified fields if not a replacement.
-            uassertStatusOK(validate(original,
-                                     noFields,
-                                     _doc,
-                                     immutablePaths,
-                                     driver->modOptions()) );
+            Status validateStatus = validate(original,
+                                             noFields,
+                                             *doc,
+                                             immutablePaths,
+                                             driver->modOptions());
+            if (!validateStatus.isOK()) {
+                return validateStatus;
+            }
         }
 
-        // Insert the doc
-        BSONObj newObj = _doc.getObject();
-        uassert(17420,
-                str::stream() << "Document to upsert is larger than " << BSONObjMaxUserSize,
-                newObj.objsize() <= BSONObjMaxUserSize);
+        BSONObj newObj = doc->getObject();
+        if (newObj.objsize() > BSONObjMaxUserSize) {
+            return Status(ErrorCodes::InvalidBSON,
+                          str::stream() << "Document to upsert is larger than "
+                                        << BSONObjMaxUserSize,
+                          17420);
+        }
+
+        *out = newObj;
+        return Status::OK();
+    }
+
+    void UpdateStage::doInsert() {
+        _specificStats.inserted = true;
+
+        const UpdateRequest* request = _params.request;
+        bool isInternalRequest = request->isFromReplication() || request->isFromMigration();
+
+        // Reset the document we will be writing to.
+        _doc.reset();
+
+        BSONObj newObj;
+        uassertStatusOK(applyUpdateOpsForInsert(_params.canonicalQuery,
+                                                request->getQuery(),
+                                                _params.driver,
+                                                request->getLifecycle(),
+                                                &_doc,
+                                                isInternalRequest,
+                                                &_specificStats,
+                                                &newObj));
 
         _specificStats.objInserted = newObj;
         if (request->shouldStoreResultDoc()) {
diff --git a/src/mongo/db/exec/update.h b/src/mongo/db/exec/update.h
index 3d820dbda6c..fdb07964bf6 100644
--- a/src/mongo/db/exec/update.h
+++ b/src/mongo/db/exec/update.h
@@ -115,6 +115,33 @@ namespace mongo {
          */
         static UpdateResult makeUpdateResult(PlanExecutor* exec, OpDebug* opDebug);
 
+        /**
+         * Computes the document to insert if the upsert flag is set to true and no matching
+         * documents are found in the database. The document to upsert is computing using the
+         * query 'cq' and the update mods contained in 'driver'.
+         *
+         * If 'cq' is NULL, which can happen for the idhack update fast path, then 'query' is
+         * used to compute the doc to insert instead of 'cq'.
+         *
+         * 'doc' is the mutable BSON document which you would like the update driver to use
+         * when computing the document to insert.
+         *
+         * Set 'isInternalRequest' to true if the upsert was issued by the replication or
+         * sharding systems.
+         *
+         * Fills out whether or not this is a fastmodinsert in 'stats'.
+         *
+         * Returns the document to insert in *out.
+         */
+        static Status applyUpdateOpsForInsert(const CanonicalQuery* cq,
+                                              const BSONObj& query,
+                                              UpdateDriver* driver,
+                                              UpdateLifecycle* lifecycle,
+                                              mutablebson::Document* doc,
+                                              bool isInternalRequest,
+                                              UpdateStats* stats,
+                                              BSONObj* out);
+
     private:
         /**
          * Computes the result of applying mods to the document 'oldObj' at RecordId 'loc' in