SERVER-50778 Compare array indexes numerically for updates

8 files changed, 97 insertions, 12 deletions

diff --git a/jstests/core/update_bit_examples.js b/jstests/core/update_bit_examples.js
index 0b8f868ea17..1f1c03596e3 100644
--- a/jstests/core/update_bit_examples.js
+++ b/jstests/core/update_bit_examples.js
@@ -1,7 +1,7 @@
 // Cannot implicitly shard accessed collections because of following errmsg: A single
 // update/delete on a sharded collection must contain an exact match on _id or contain the shard
 // key.
-// @tags: [assumes_unsharded_collection, requires_non_retryable_writes]
+// @tags: [assumes_unsharded_collection, requires_non_retryable_writes, requires_fcv_47]
 
 // Basic examples for $bit
 var res;
@@ -32,3 +32,10 @@ assert.eq(coll.findOne().a, 4);
 // SERVER-19706 Empty bit operation.
 res = coll.update({}, {$bit: {a: {}}});
 assert.writeError(res);
+
+// Make sure $bit on index arrays 9 and 10 when padding is needed works.
+assert.commandWorked(coll.insert({_id: 2, a: [0]}));
+assert.commandWorked(
+    coll.update({_id: 2}, {$bit: {"a.9": {or: NumberInt(0)}, "a.10": {or: NumberInt(0)}}}));
+res = coll.find({_id: 2}).toArray();
+assert.eq(res[0]["a"], [0, null, null, null, null, null, null, null, null, 0, 0]);
diff --git a/jstests/multiVersion/bit_update_mixed_fcv.js b/jstests/multiVersion/bit_update_mixed_fcv.js
new file mode 100644
index 00000000000..f3012a6d382
--- /dev/null
+++ b/jstests/multiVersion/bit_update_mixed_fcv.js
@@ -0,0 +1,29 @@
+/**
+ * Tests that $bit updates succeed if primaries compare updates numerically and secondaries compare
+ * updates lexicographically for array indexes.
+ */
+
+(function() {
+"use strict";
+
+const rst = new ReplSetTest({
+    nodes: 2,
+    nodeOptions: {binVersion: "latest"},
+});
+
+rst.startSet();
+rst.initiate();
+
+const primary = rst.getPrimary();
+const testDB = primary.getDB(jsTestName());
+const coll = testDB.test;
+
+testDB.adminCommand({setFeatureCompatibilityVersion: lastLTSFCV});
+
+assert.commandWorked(coll.insert({_id: 0, arr: [0]}));
+
+assert.commandWorked(
+    coll.update({_id: 0}, {$bit: {"a.9": {or: NumberInt(0)}, "a.10": {or: NumberInt(0)}}}));
+
+rst.stopSet();
+})();
diff --git a/src/mongo/db/update/path_support.h b/src/mongo/db/update/path_support.h
index 808f2df5dff..5698bffd7ef 100644
--- a/src/mongo/db/update/path_support.h
+++ b/src/mongo/db/update/path_support.h
@@ -38,6 +38,7 @@
 #include "mongo/db/field_ref_set.h"
 #include "mongo/db/matcher/expression.h"
 #include "mongo/db/matcher/expression_leaf.h"
+#include "mongo/util/ctype.h"
 
 namespace mongo {
 
@@ -50,6 +51,46 @@ static const size_t kMaxPaddingAllowed = 1500000;
 // Convenience type to hold equality matches at particular paths from a MatchExpression
 typedef std::map<StringData, const EqualityMatchExpression*> EqualityMatches;
 
+struct cmpPathsAndArrayIndexes {
+    // While there is a string to number parser in the codebase, we use this for performance
+    // reasons. This code allows us to only look at the characters/digits that are relevant for our
+    // comparisons.
+    bool operator()(const std::string& a, const std::string& b) const {
+        auto aLength = a.size();
+        auto bLength = b.size();
+        // If these are not numbers, do a normal string comparison. Treat zero padded numbers as
+        // strings.
+        if (aLength == 0 || bLength == 0 || !ctype::isDigit(a[0]) || !ctype::isDigit(b[0]) ||
+            (a[0] == '0' && aLength > 1) || (b[0] == '0' && bLength > 1)) {
+            return a < b;
+        }
+
+        // At this point we are only seeing strings that are prefixed with at least one digit. If
+        // they are numbers, the longer one is bigger because it has a larger order of magnitude.
+        // If they are not numbers (ex: 123qyz) we would not be able to distinguish which is larger
+        // without inspecting every character. To avoid this work we choose to order by length
+        // because the actual order is not important but needs to be consistent.
+        if (aLength != bLength) {
+            return aLength < bLength;
+        }
+        // Find the larger number. Return immediately if a high order digit doesnt match, or we find
+        // a non-number.
+        size_t aIndex = 0;
+        size_t bIndex = 0;
+        while (ctype::isDigit(a[aIndex]) && ctype::isDigit(b[bIndex]) && aIndex < aLength &&
+               bIndex < bLength) {
+            if (a[aIndex] == b[bIndex]) {
+                ++aIndex;
+                ++bIndex;
+                continue;
+            }
+            return a[aIndex] < b[bIndex];
+        }
+        // We found a non-number, perform normal string comparison.
+        return a < b;
+    }
+};
+
 /**
  * Finds the longest portion of 'prefix' that exists in document rooted at 'root' and is
  * "viable." A viable path is one that, if fully created on a given doc, would not
diff --git a/src/mongo/db/update/update_array_node.h b/src/mongo/db/update/update_array_node.h
index c6e90c1d9c3..ff48b8dba59 100644
--- a/src/mongo/db/update/update_array_node.h
+++ b/src/mongo/db/update/update_array_node.h
@@ -99,7 +99,7 @@ public:
 
 private:
     const std::map<StringData, std::unique_ptr<ExpressionWithPlaceholder>>& _arrayFilters;
-    std::map<std::string, clonable_ptr<UpdateNode>> _children;
+    std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes> _children;
 
     // When calling apply() causes us to merge elements of '_children', we store the result of the
     // merge in case we need it for another array element or document.
diff --git a/src/mongo/db/update/update_internal_node.cpp b/src/mongo/db/update/update_internal_node.cpp
index 95cc60eda8e..fe5e14ab336 100644
--- a/src/mongo/db/update/update_internal_node.cpp
+++ b/src/mongo/db/update/update_internal_node.cpp
@@ -34,12 +34,15 @@
 namespace mongo {
 
 // static
-std::map<std::string, clonable_ptr<UpdateNode>> UpdateInternalNode::createUpdateNodeMapByMerging(
-    const std::map<std::string, clonable_ptr<UpdateNode>>& leftMap,
-    const std::map<std::string, clonable_ptr<UpdateNode>>& rightMap,
+std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>
+UpdateInternalNode::createUpdateNodeMapByMerging(
+    const std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>&
+        leftMap,
+    const std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>&
+        rightMap,
     FieldRef* pathTaken,
     bool wrapFieldNameAsArrayFilterIdentifier) {
-    std::map<std::string, clonable_ptr<UpdateNode>> mergedMap;
+    std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes> mergedMap;
 
     // Get the union of the field names we know about among the leftMap and rightMap.
     stdx::unordered_set<std::string> allFields;
diff --git a/src/mongo/db/update/update_internal_node.h b/src/mongo/db/update/update_internal_node.h
index 9cadd04fc57..a3403879513 100644
--- a/src/mongo/db/update/update_internal_node.h
+++ b/src/mongo/db/update/update_internal_node.h
@@ -34,6 +34,7 @@
 
 #include "mongo/base/clonable_ptr.h"
 #include "mongo/db/field_ref.h"
+#include "mongo/db/update/path_support.h"
 #include "mongo/db/update/update_node.h"
 
 namespace mongo {
@@ -72,9 +73,12 @@ protected:
      * wrapFieldNameAsArrayFilterIdentifier is true, field names are wrapped as $[<field name>] for
      * error reporting.
      */
-    static std::map<std::string, clonable_ptr<UpdateNode>> createUpdateNodeMapByMerging(
-        const std::map<std::string, clonable_ptr<UpdateNode>>& leftMap,
-        const std::map<std::string, clonable_ptr<UpdateNode>>& rightMap,
+    static std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>
+    createUpdateNodeMapByMerging(
+        const std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>&
+            leftMap,
+        const std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes>&
+            rightMap,
         FieldRef* pathTaken,
         bool wrapFieldNameAsArrayFilterIdentifier = false);
 
diff --git a/src/mongo/db/update/update_object_node.h b/src/mongo/db/update/update_object_node.h
index d7f2e56e9de..8c347968e74 100644
--- a/src/mongo/db/update/update_object_node.h
+++ b/src/mongo/db/update/update_object_node.h
@@ -39,6 +39,7 @@
 #include "mongo/bson/bsonelement.h"
 #include "mongo/db/matcher/expression_with_placeholder.h"
 #include "mongo/db/update/modifier_table.h"
+#include "mongo/db/update/path_support.h"
 #include "mongo/db/update/update_internal_node.h"
 #include "mongo/stdx/unordered_map.h"
 
@@ -126,12 +127,12 @@ public:
         visitor->visit(this);
     }
 
-    const std::map<std::string, clonable_ptr<UpdateNode>>& getChildren() const {
+    const auto& getChildren() const {
         return _children;
     }
 
 private:
-    std::map<std::string, clonable_ptr<UpdateNode>> _children;
+    std::map<std::string, clonable_ptr<UpdateNode>, pathsupport::cmpPathsAndArrayIndexes> _children;
     clonable_ptr<UpdateNode> _positionalChild;
 
     // When calling apply() causes us to merge an element of '_children' with '_positionalChild', we
diff --git a/src/mongo/db/update/update_object_node_test.cpp b/src/mongo/db/update/update_object_node_test.cpp
index a8d09f4f891..2e5906e8f30 100644
--- a/src/mongo/db/update/update_object_node_test.cpp
+++ b/src/mongo/db/update/update_object_node_test.cpp
@@ -2565,7 +2565,7 @@ TEST_F(UpdateObjectNodeTest, ApplyMultipleArrayUpdates) {
         doc.getObject());
     ASSERT_FALSE(doc.isInPlaceModeEnabled());
 
-    assertOplogEntry(fromjson("{$set: {'a.10': 10, 'a.2': 2}}"),
+    assertOplogEntry(fromjson("{$set: {'a.2': 2, 'a.10': 10}}"),
                      fromjson("{$v: 2, diff: {sa: {a: true, u2: 2, u10: 10}}}"));
 }