SERVER-37124 Retry full upsert path when duplicate key exception matches exact query predicate

2 files changed, 110 insertions, 0 deletions

diff --git a/jstests/concurrency/fsm_workloads/upsert_unique_index.js b/jstests/concurrency/fsm_workloads/upsert_unique_index.js
new file mode 100644
index 00000000000..62d91794ccb
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/upsert_unique_index.js
@@ -0,0 +1,44 @@
+'use strict';
+
+/**
+ * Performs concurrent upsert and delete operations against a small set of documents with a unique
+ * index in place. One specific scenario this test exercises is upsert retry in the case where an
+ * upsert generates an insert, which then fails due to another operation inserting first.
+ */
+var $config = (function() {
+
+    const data = {
+        numDocs: 4,
+        getDocValue: function() {
+            return Random.randInt(this.numDocs);
+        },
+    };
+
+    const states = {
+        delete: function(db, collName) {
+            assertAlways.commandWorked(
+                db[collName].remove({_id: this.getDocValue()}, {justOne: true}));
+        },
+        upsert: function(db, collName) {
+            const value = this.getDocValue();
+            const cmdRes = db.runCommand(
+                {update: collName, updates: [{q: {_id: value}, u: {$inc: {y: 1}}, upsert: true}]});
+
+            assertAlways.commandWorked(cmdRes);
+        },
+    };
+
+    const transitions = {
+        upsert: {upsert: 0.5, delete: 0.5},
+        delete: {upsert: 0.5, delete: 0.5},
+    };
+
+    return {
+        threadCount: 20,
+        iterations: 100,
+        states: states,
+        startState: 'upsert',
+        transitions: transitions,
+        data: data,
+    };
+})();
diff --git a/jstests/noPassthrough/upsert_duplicate_key_retry.js b/jstests/noPassthrough/upsert_duplicate_key_retry.js
new file mode 100644
index 00000000000..c35e490377f
--- /dev/null
+++ b/jstests/noPassthrough/upsert_duplicate_key_retry.js
@@ -0,0 +1,66 @@
+/**
+ * When two concurrent identical upsert operations are performed, for which a unique index exists on
+ * the query values, it is possible that they will both attempt to perform an insert with one of
+ * the two failing on the unique index constraint. This test confirms that the failed insert will be
+ * retried, resulting in an update.
+ *  @tags: [requires_replication]
+ */
+
+(function() {
+    "use strict";
+
+    const rst = new ReplSetTest({nodes: 1});
+    rst.startSet();
+    rst.initiate();
+
+    const testDB = rst.getPrimary().getDB("test");
+    const adminDB = testDB.getSiblingDB("admin");
+    const collName = "upsert_duplicate_key_retry";
+    const testColl = testDB.getCollection(collName);
+
+    testDB.runCommand({drop: collName});
+
+    // Queries current operations until 'count' matching operations are found.
+    function awaitMatchingCurrentOpCount(message, count) {
+        assert.soon(() => {
+            const currentOp =
+                adminDB.aggregate([{$currentOp: {}}, {$match: {msg: message}}]).toArray();
+            return (currentOp.length === count);
+        });
+    }
+
+    function performUpsert() {
+        // This function is called from startParallelShell(), so closed-over variables will not be
+        // available. We must re-obtain the value of 'testColl' in the function body.
+        const testColl = db.getMongo().getDB("test").getCollection("upsert_duplicate_key_retry");
+        assert.commandWorked(testColl.update({x: 3}, {$inc: {y: 1}}, {upsert: true}));
+    }
+
+    assert.commandWorked(testColl.createIndex({x: 1}, {unique: true}));
+
+    // Will hang upsert operations just prior to performing an insert.
+    assert.commandWorked(testDB.adminCommand(
+        {configureFailPoint: "hangBeforeUpsertPerformsInsert", mode: "alwaysOn"}));
+
+    const awaitUpdate1 = startParallelShell(performUpsert, rst.ports[0]);
+    const awaitUpdate2 = startParallelShell(performUpsert, rst.ports[0]);
+
+    awaitMatchingCurrentOpCount("hangBeforeUpsertPerformsInsert", 2);
+
+    assert.commandWorked(
+        testDB.adminCommand({configureFailPoint: "hangBeforeUpsertPerformsInsert", mode: "off"}));
+
+    awaitUpdate1();
+    awaitUpdate2();
+
+    const cursor = testColl.find({}, {_id: 0});
+    assert.eq(cursor.next(), {x: 3, y: 2});
+    assert(!cursor.hasNext(), cursor.toArray());
+
+    // Confirm that oplog entries exist for both insert and update operation.
+    const oplogColl = testDB.getSiblingDB("local").getCollection("oplog.rs");
+    assert.eq(1, oplogColl.find({"op": "i", "ns": "test.upsert_duplicate_key_retry"}).itcount());
+    assert.eq(1, oplogColl.find({"op": "u", "ns": "test.upsert_duplicate_key_retry"}).itcount());
+
+    rst.stopSet();
+})();