SERVER-38461 Limit upsert retry to DuplicateKey violations with matching values

6 files changed, 35 insertions, 10 deletions

diff --git a/src/mongo/db/exec/update.cpp b/src/mongo/db/exec/update.cpp
index 80749d64f8f..60e081322ed 100644
--- a/src/mongo/db/exec/update.cpp
+++ b/src/mongo/db/exec/update.cpp
@@ -37,6 +37,7 @@
 #include <algorithm>
 
 #include "mongo/base/status_with.h"
+#include "mongo/bson/bson_comparator_interface_base.h"
 #include "mongo/bson/mutable/algorithm.h"
 #include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/curop_failpoint_helpers.h"
@@ -476,7 +477,8 @@ bool UpdateStage::shouldRetryDuplicateKeyException(const ParsedUpdate& parsedUpd
     }
 
     // In order to be retryable, the update equality field paths must be identical to the unique
-    // index key field paths.
+    // index key field paths. Also, the values that triggered the DuplicateKey error must match the
+    // values used in the upsert query predicate.
     pathsupport::EqualityMatches equalities;
     auto status = pathsupport::extractEqualityMatches(*matchExpr, &equalities);
     if (!status.isOK()) {
@@ -488,11 +490,27 @@ bool UpdateStage::shouldRetryDuplicateKeyException(const ParsedUpdate& parsedUpd
         return false;
     }
 
-    for (const auto& key : keyPattern) {
-        if (!equalities.count(key.fieldNameStringData())) {
+    auto keyValue = errorInfo.getDuplicatedKeyValue();
+
+    BSONObjIterator keyPatternIter(keyPattern);
+    BSONObjIterator keyValueIter(keyValue);
+    while (keyPatternIter.more() && keyValueIter.more()) {
+        auto keyPatternElem = keyPatternIter.next();
+        auto keyValueElem = keyValueIter.next();
+
+        auto keyName = keyPatternElem.fieldNameStringData();
+        if (!equalities.count(keyName)) {
+            return false;
+        }
+
+        // Comparison which obeys field ordering but ignores field name.
+        BSONElementComparator cmp{BSONElementComparator::FieldNamesMode::kIgnore, nullptr};
+        if (cmp.evaluate(equalities[keyName]->getData() != keyValueElem)) {
             return false;
         }
     }
+    invariant(!keyPatternIter.more());
+    invariant(!keyValueIter.more());
 
     return true;
 }
diff --git a/src/mongo/db/storage/duplicate_key_error_info.cpp b/src/mongo/db/storage/duplicate_key_error_info.cpp
index 911b3de0381..deeb075898f 100644
--- a/src/mongo/db/storage/duplicate_key_error_info.cpp
+++ b/src/mongo/db/storage/duplicate_key_error_info.cpp
@@ -44,10 +44,11 @@ MONGO_INIT_REGISTER_ERROR_EXTRA_INFO(DuplicateKeyErrorInfo);
 
 void DuplicateKeyErrorInfo::serialize(BSONObjBuilder* bob) const {
     bob->append("keyPattern", _keyPattern);
+    bob->append("keyValue", _keyValue);
 }
 
 std::shared_ptr<const ErrorExtraInfo> DuplicateKeyErrorInfo::parse(const BSONObj& obj) {
-    return std::make_shared<DuplicateKeyErrorInfo>(obj["keyPattern"].Obj());
+    return std::make_shared<DuplicateKeyErrorInfo>(obj["keyPattern"].Obj(), obj["keyValue"].Obj());
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/duplicate_key_error_info.h b/src/mongo/db/storage/duplicate_key_error_info.h
index 06a17cce94c..d913412f5d9 100644
--- a/src/mongo/db/storage/duplicate_key_error_info.h
+++ b/src/mongo/db/storage/duplicate_key_error_info.h
@@ -45,8 +45,8 @@ public:
 
     static std::shared_ptr<const ErrorExtraInfo> parse(const BSONObj&);
 
-    explicit DuplicateKeyErrorInfo(const BSONObj& keyPattern)
-        : _keyPattern(keyPattern.getOwned()) {}
+    explicit DuplicateKeyErrorInfo(const BSONObj& keyPattern, const BSONObj& keyValue)
+        : _keyPattern(keyPattern.getOwned()), _keyValue(keyValue.getOwned()) {}
 
     void serialize(BSONObjBuilder* bob) const override;
 
@@ -60,8 +60,13 @@ public:
         return _keyPattern;
     }
 
+    const BSONObj& getDuplicatedKeyValue() const {
+        return _keyValue;
+    }
+
 private:
-    const BSONObj _keyPattern;
+    BSONObj _keyPattern;
+    BSONObj _keyValue;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/storage/index_entry_comparison.cpp b/src/mongo/db/storage/index_entry_comparison.cpp
index 67613639633..d597cfe082c 100644
--- a/src/mongo/db/storage/index_entry_comparison.cpp
+++ b/src/mongo/db/storage/index_entry_comparison.cpp
@@ -195,7 +195,7 @@ Status buildDupKeyErrorStatus(const BSONObj& key,
     }
 
     sb << builder.obj();
-    return Status(DuplicateKeyErrorInfo(keyPattern), sb.str());
+    return Status(DuplicateKeyErrorInfo(keyPattern, key), sb.str());
 }
 
 }  // namespace mongo
diff --git a/src/mongo/s/catalog/dist_lock_catalog_impl_test.cpp b/src/mongo/s/catalog/dist_lock_catalog_impl_test.cpp
index 8eac156a8ba..12e6dff6bbd 100644
--- a/src/mongo/s/catalog/dist_lock_catalog_impl_test.cpp
+++ b/src/mongo/s/catalog/dist_lock_catalog_impl_test.cpp
@@ -438,7 +438,8 @@ TEST_F(DistLockCatalogTest, GrabLockDupKeyError) {
     });
 
     onCommand([](const RemoteCommandRequest& request) -> StatusWith<BSONObj> {
-        return Status({DuplicateKeyErrorInfo(BSON("x" << 1)), "Mock duplicate key error"});
+        return Status(
+            {DuplicateKeyErrorInfo(BSON("x" << 1), BSON("" << 1)), "Mock duplicate key error"});
     });
 
     future.timed_get(kFutureTimeout);
diff --git a/src/mongo/s/catalog/sharding_catalog_write_retry_test.cpp b/src/mongo/s/catalog/sharding_catalog_write_retry_test.cpp
index 7ce5acce70c..13d86abed00 100644
--- a/src/mongo/s/catalog/sharding_catalog_write_retry_test.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_write_retry_test.cpp
@@ -80,7 +80,7 @@ const HostAndPort kTestHosts[] = {
     HostAndPort("TestHost1:12345"), HostAndPort("TestHost2:12345"), HostAndPort("TestHost3:12345")};
 
 Status getMockDuplicateKeyError() {
-    return {DuplicateKeyErrorInfo(BSON("mock" << 1)), "Mock duplicate key error"};
+    return {DuplicateKeyErrorInfo(BSON("mock" << 1), BSON("" << 1)), "Mock duplicate key error"};
 }
 
 TEST_F(InsertRetryTest, RetryOnInterruptedAndNetworkErrorSuccess) {