SERVER-30302 Add explicit configuration object for UpdateSequenceGenerator

3 files changed, 158 insertions, 47 deletions

diff --git a/src/mongo/db/repl/idempotency_update_sequence.cpp b/src/mongo/db/repl/idempotency_update_sequence.cpp
index b235256e846..7ed6ea23167 100644
--- a/src/mongo/db/repl/idempotency_update_sequence.cpp
+++ b/src/mongo/db/repl/idempotency_update_sequence.cpp
@@ -39,6 +39,19 @@
 
 namespace mongo {
 
+UpdateSequenceGeneratorConfig::UpdateSequenceGeneratorConfig(std::set<StringData> fields_,
+                                                             size_t depth_,
+                                                             size_t length_,
+                                                             double scalarProbability_,
+                                                             double docProbability_,
+                                                             double arrProbability_)
+    : fields(std::move(fields_)),
+      depth(depth_),
+      length(length_),
+      scalarProbability(scalarProbability_),
+      docProbability(docProbability_),
+      arrProbability(arrProbability_) {}
+
 std::size_t UpdateSequenceGenerator::_getPathDepth(const std::string& path) {
     // Our depth is -1 because we count at 0, but numParts just counts the number of fields.
     return path == "" ? 0 : FieldRef(path).numParts() - 1;
@@ -57,36 +70,40 @@ std::vector<std::string> UpdateSequenceGenerator::_eliminatePrefixPaths(
     return remainingPaths;
 }
 
-void UpdateSequenceGenerator::_generatePaths(const std::set<StringData>& fields,
-                                             const std::size_t depth,
-                                             const std::size_t length,
+void UpdateSequenceGenerator::_generatePaths(const UpdateSequenceGeneratorConfig& config,
                                              const std::string& path) {
-    if (UpdateSequenceGenerator::_getPathDepth(path) == depth) {
+    if (UpdateSequenceGenerator::_getPathDepth(path) == config.depth) {
         return;
     }
 
     if (!path.empty()) {
-        for (std::size_t i = 0; i < length; i++) {
+        for (std::size_t i = 0; i < config.length; i++) {
             FieldRef arrPathRef(path);
             arrPathRef.appendPart(std::to_string(i));
             auto arrPath = arrPathRef.dottedField().toString();
             _paths.push_back(arrPath);
-            _generatePaths(fields, depth, length, arrPath);
+            _generatePaths(config, arrPath);
         }
     }
 
-    if (fields.empty()) {
+    if (config.fields.empty()) {
         return;
     }
 
-    std::set<StringData> remainingFields(fields);
-    for (auto field : fields) {
+    std::set<StringData> remainingFields(config.fields);
+    for (auto field : config.fields) {
         remainingFields.erase(remainingFields.begin());
         FieldRef docPathRef(path);
         docPathRef.appendPart(field);
         auto docPath = docPathRef.dottedField().toString();
         _paths.push_back(docPath);
-        _generatePaths(remainingFields, depth, length, docPath);
+        UpdateSequenceGeneratorConfig remainingConfig = {remainingFields,
+                                                         config.depth,
+                                                         config.length,
+                                                         config.scalarProbability,
+                                                         config.arrProbability,
+                                                         config.docProbability};
+        _generatePaths(remainingConfig, docPath);
     }
 }
 
@@ -109,9 +126,10 @@ std::vector<std::string> UpdateSequenceGenerator::_getRandomPaths() const {
 }
 
 BSONObj UpdateSequenceGenerator::generateUpdate() const {
-    bool generateSetUpdate =
-        this->_random.nextInt32(UpdateSequenceGenerator::kNumUpdateChoices) == 1;
-    if (generateSetUpdate) {
+    double setSum = this->_config.scalarProbability + this->_config.arrProbability +
+        this->_config.docProbability;
+    double generateSetUpdate = this->_random.nextCanonicalDouble();
+    if (generateSetUpdate <= setSum) {
         return _generateSet();
     } else {
         return _generateUnset();
@@ -132,12 +150,22 @@ BSONObj UpdateSequenceGenerator::_generateSet() const {
 
 UpdateSequenceGenerator::SetChoice UpdateSequenceGenerator::_determineWhatToSet(
     const std::string& setPath) const {
-    if (UpdateSequenceGenerator::_getPathDepth(setPath) == this->_depth) {
+    if (UpdateSequenceGenerator::_getPathDepth(setPath) == this->_config.depth) {
+        // If we have hit the max depth, we don't have a choice anyways.
         auto choice = static_cast<size_t>(SetChoice::kNumScalarSetChoices);
         return static_cast<SetChoice>(this->_random.nextInt32(choice));
     } else {
-        auto choice = static_cast<size_t>(SetChoice::kNumScalarSetChoices);
-        return static_cast<SetChoice>(UpdateSequenceGenerator::_random.nextInt32(choice));
+        double setSum = this->_config.scalarProbability + this->_config.arrProbability +
+            this->_config.docProbability;
+        double choice = this->_random.nextCanonicalDouble() * setSum;
+        if (choice <= this->_config.scalarProbability) {
+            auto scalarChoice = static_cast<size_t>(SetChoice::kNumScalarSetChoices);
+            return static_cast<SetChoice>(this->_random.nextInt32(scalarChoice));
+        } else if (choice <= setSum - this->_config.docProbability) {
+            return SetChoice::kSetArr;
+        } else {
+            return SetChoice::kSetDoc;
+        }
     }
 }
 
@@ -204,7 +232,7 @@ BSONObj UpdateSequenceGenerator::_generateDocToSet(const std::string& setPath) c
 }
 
 std::set<StringData> UpdateSequenceGenerator::_getRemainingFields(const std::string& path) const {
-    std::set<StringData> remainingFields(this->_fields);
+    std::set<StringData> remainingFields(this->_config.fields);
 
     FieldRef pathRef(path);
     StringData lastField;
@@ -212,7 +240,7 @@ std::set<StringData> UpdateSequenceGenerator::_getRemainingFields(const std::str
     // array positions (numbers).
     for (int i = pathRef.numParts() - 1; i >= 0; i--) {
         auto field = pathRef.getPart(i);
-        if (this->_fields.find(field) != this->_fields.end()) {
+        if (this->_config.fields.find(field) != this->_config.fields.end()) {
             lastField = field;
             break;
         }
@@ -221,7 +249,7 @@ std::set<StringData> UpdateSequenceGenerator::_getRemainingFields(const std::str
     // The last alphabetic field used must be after all other alphabetic fields that could ever be
     // used, since the fields that are used are selected in the order that they pop off from a
     // std::set.
-    for (auto field : this->_fields) {
+    for (auto field : this->_config.fields) {
         remainingFields.erase(field);
         if (field == lastField) {
             break;
@@ -234,12 +262,12 @@ std::set<StringData> UpdateSequenceGenerator::_getRemainingFields(const std::str
 DocumentStructureEnumerator UpdateSequenceGenerator::_getValidEnumeratorForPath(
     const std::string& path) const {
     auto remainingFields = _getRemainingFields(path);
-    std::size_t remainingDepth = this->_depth - UpdateSequenceGenerator::_getPathDepth(path);
+    std::size_t remainingDepth = this->_config.depth - UpdateSequenceGenerator::_getPathDepth(path);
     if (remainingDepth > 0) {
         remainingDepth -= 1;
     }
 
-    DocumentStructureEnumerator enumerator({remainingFields, remainingDepth, this->_length});
+    DocumentStructureEnumerator enumerator({remainingFields, remainingDepth, this->_config.length});
     return enumerator;
 }
 
@@ -251,15 +279,11 @@ std::vector<std::string> UpdateSequenceGenerator::getPaths() const {
     return this->_paths;
 }
 
-UpdateSequenceGenerator::UpdateSequenceGenerator(std::set<StringData> fields,
-                                                 std::size_t depth,
-                                                 std::size_t length)
-    : _fields(fields),
-      _depth(depth),
-      _length(length),
+UpdateSequenceGenerator::UpdateSequenceGenerator(UpdateSequenceGeneratorConfig config)
+    : _config(std::move(config)),
       _random(std::unique_ptr<SecureRandom>(SecureRandom::create())->nextInt64()) {
     auto path = "";
-    _generatePaths(fields, depth, length, path);
+    _generatePaths(config, path);
     // Creates the same shuffle each time, but we don't care. We want to mess up the DFS ordering.
     std::random_shuffle(this->_paths.begin(), this->_paths.end(), this->_random);
 }
diff --git a/src/mongo/db/repl/idempotency_update_sequence.h b/src/mongo/db/repl/idempotency_update_sequence.h
index d31360c3a9d..e83b0553147 100644
--- a/src/mongo/db/repl/idempotency_update_sequence.h
+++ b/src/mongo/db/repl/idempotency_update_sequence.h
@@ -44,10 +44,26 @@ class BSONObj;
 struct BSONArray;
 class BSONObjBuilder;
 
+struct UpdateSequenceGeneratorConfig {
+    UpdateSequenceGeneratorConfig(std::set<StringData> fields_,
+                                  size_t depth_,
+                                  size_t length_,
+                                  double scalarProbability_ = 0.250,
+                                  double docProbability_ = 0.250,
+                                  double arrProbability_ = 0.250);
+
+    const std::set<StringData> fields = {};
+    const size_t depth = 0;
+    const size_t length = 0;
+    const double scalarProbability = 0.250;
+    const double docProbability = 0.250;
+    const double arrProbability = 0.250;
+};
+
 class UpdateSequenceGenerator : public SequenceGenerator {
 
 public:
-    UpdateSequenceGenerator(std::set<StringData> fields, std::size_t depth, std::size_t length);
+    explicit UpdateSequenceGenerator(UpdateSequenceGeneratorConfig config);
 
     BSONObj generateUpdate() const;
 
@@ -88,10 +104,7 @@ private:
     static std::vector<std::string> _eliminatePrefixPaths(const std::string& path,
                                                           const std::vector<std::string>& paths);
 
-    void _generatePaths(const std::set<StringData>& fields,
-                        const std::size_t depth,
-                        const std::size_t length,
-                        const std::string& path);
+    void _generatePaths(const UpdateSequenceGeneratorConfig& config, const std::string& path);
 
     std::set<StringData> _getRemainingFields(const std::string& path) const;
 
@@ -116,9 +129,7 @@ private:
     BSONObj _generateDocToSet(const std::string& setPath) const;
 
     std::vector<std::string> _paths;
-    const std::set<StringData> _fields;
-    const std::size_t _depth;
-    const std::size_t _length;
+    const UpdateSequenceGeneratorConfig _config;
     mutable PseudoRandom _random;
 };
 
diff --git a/src/mongo/db/repl/idempotency_update_sequence_test.cpp b/src/mongo/db/repl/idempotency_update_sequence_test.cpp
index bb8d930158d..3d5edf11700 100644
--- a/src/mongo/db/repl/idempotency_update_sequence_test.cpp
+++ b/src/mongo/db/repl/idempotency_update_sequence_test.cpp
@@ -55,7 +55,7 @@ TEST(UpdateGenTest, FindsAllPaths) {
     std::set<StringData> fields{"a", "b"};
     size_t depth = 1;
     size_t length = 1;
-    UpdateSequenceGenerator generator(fields, depth, length);
+    UpdateSequenceGenerator generator({fields, depth, length});
 
     ASSERT_EQ(generator.getPaths().size(), 5U);
 
@@ -84,7 +84,7 @@ TEST(UpdateGenTest, NoDuplicatePaths) {
     std::set<StringData> fields{"a", "b"};
     size_t depth = 2;
     size_t length = 2;
-    UpdateSequenceGenerator generator(fields, depth, length);
+    UpdateSequenceGenerator generator({fields, depth, length});
 
     auto paths = generator.getPaths();
     for (size_t i = 0; i < paths.size(); i++) {
@@ -103,7 +103,7 @@ TEST(UpdateGenTest, UpdatesHaveValidPaths) {
     std::set<StringData> fields{"a", "b"};
     size_t depth = 1;
     size_t length = 1;
-    UpdateSequenceGenerator generator(fields, depth, length);
+    UpdateSequenceGenerator generator({fields, depth, length});
     auto update = generator.generateUpdate();
 
     BSONObj updateArg;
@@ -140,7 +140,7 @@ TEST(UpdateGenTest, UpdatesAreNotAmbiguous) {
     std::set<StringData> fields{"a", "b"};
     size_t depth = 1;
     size_t length = 1;
-    UpdateSequenceGenerator generator(fields, depth, length);
+    UpdateSequenceGenerator generator({fields, depth, length});
     auto update = generator.generateUpdate();
 
     BSONObj updateArg;
@@ -180,21 +180,19 @@ std::size_t getMaxDepth(BSONObj obj) {
     }
 
     return curMaxDepth;
-};
+}
 
 TEST(UpdateGenTest, UpdatesPreserveDepthConstraint) {
     std::set<StringData> fields{"a", "b"};
     size_t depth = 2;
     size_t length = 1;
-    UpdateSequenceGenerator generator(fields, depth, length);
+    UpdateSequenceGenerator generator({fields, depth, length, 0.333, 0.333, 0.334});
 
     BSONElement setElem;
     BSONObj update;
-    // Keep trying until we get a $set, so that we can avoid a no-op test case.
-    while (!setElem) {
-        update = generator.generateUpdate();
-        setElem = update["$set"];
-    }
+    // Because our probabilities sum to 1, we are guaranteed to always get a $set.
+    update = generator.generateUpdate();
+    setElem = update["$set"];
     BSONObj updateArg = setElem.Obj();
 
     std::set<std::string> argPaths;
@@ -217,5 +215,83 @@ TEST(UpdateGenTest, UpdatesPreserveDepthConstraint) {
         }
     }
 }
+
+TEST(UpdateGenTest, OnlyGenerateUnset) {
+    std::set<StringData> fields{"a", "b"};
+    size_t depth = 1;
+    size_t length = 1;
+
+    UpdateSequenceGenerator generatorNoSet({fields, depth, length, 0.0, 0.0, 0.0});
+    for (size_t i = 0; i < 100; i++) {
+        auto update = generatorNoSet.generateUpdate();
+        if (!update["$unset"]) {
+            StringBuilder sb;
+            sb << "Generator created an update that was not an $unset, even though the probability "
+                  "of doing so is zero: "
+               << update;
+            FAIL(sb.str());
+        }
+    }
+}
+
+TEST(UpdateGenTest, OnlySetUpdatesWithScalarValue) {
+    std::set<StringData> fields{"a", "b"};
+    size_t depth = 1;
+    size_t length = 1;
+    UpdateSequenceGenerator generatorNoUnsetAndOnlyScalar({fields, depth, length, 1.0, 0.0, 0.0});
+    for (size_t i = 0; i < 100; i++) {
+        auto update = generatorNoUnsetAndOnlyScalar.generateUpdate();
+        if (!update["$set"]) {
+            StringBuilder sb;
+            sb << "Generator created an update that was not an $set, even though the probability "
+                  "of doing so is zero: "
+               << update;
+            FAIL(sb.str());
+        } else if (getMaxDepth(update["$set"].Obj()) != 0) {
+            StringBuilder sb;
+            sb << "Generator created an update that had a nonscalar value, because it's maximum "
+                  "depth was nonzero: "
+               << update;
+            FAIL(sb.str());
+        }
+    }
+}
+
+TEST(UpdateGenTest, OnlySetUpdatesWithScalarsAtMaxDepth) {
+    std::set<StringData> fields{"a", "b"};
+    size_t depth = 2;
+    size_t length = 1;
+    UpdateSequenceGenerator generatorNeverScalar({fields, depth, length, 0.0, 0.5, 0.5});
+    for (size_t i = 0; i < 100; i++) {
+        auto update = generatorNeverScalar.generateUpdate();
+        for (auto elem : update["$set"].Obj()) {
+            StringData fieldName = elem.fieldNameStringData();
+            FieldRef fieldRef(fieldName);
+            size_t pathDepth = getPathDepth_forTest(fieldName.toString());
+            bool isDocOrArr = elem.type() == BSONType::Object || elem.type() == BSONType::Array;
+            if (pathDepth != depth) {
+                // If the path is not equal to the max depth we provided above, then there
+                // should
+                // only be an array or doc at this point.
+                if (!isDocOrArr) {
+                    StringBuilder sb;
+                    sb << "The set argument: " << elem
+                       << " is a scalar, but the probability of a scalar occuring for a path that "
+                          "does not meet the maximum depth is zero.";
+                    FAIL(sb.str());
+                }
+            } else {
+                if (isDocOrArr) {
+                    StringBuilder sb;
+                    sb << "The set argument: " << elem
+                       << " is not scalar, however, this path reaches the maximum depth so a "
+                          "scalar should be the only choice.";
+                    FAIL(sb.str());
+                }
+            }
+        }
+    }
+}
+
 }  // namespace
 }  // namespace mongo