SERVER-13492 add OID and SEQ_INT to bson template evaluator

5 files changed, 302 insertions, 23 deletions

diff --git a/src/mongo/scripting/bench.cpp b/src/mongo/scripting/bench.cpp
index b85d0b3c5a7..e2fee9b12b2 100644
--- a/src/mongo/scripting/bench.cpp
+++ b/src/mongo/scripting/bench.cpp
@@ -308,8 +308,8 @@ namespace mongo {
         return b.obj();
     }
 
-    BenchRunWorker::BenchRunWorker(const BenchRunConfig *config, BenchRunState *brState)
-        : _config(config), _brState(brState) {
+    BenchRunWorker::BenchRunWorker(size_t id, const BenchRunConfig *config, BenchRunState *brState)
+        : _id(id), _config(config), _brState(brState) {
     }
 
     BenchRunWorker::~BenchRunWorker() {}
@@ -330,6 +330,7 @@ namespace mongo {
         mongo::Timer timer;
 
         BsonTemplateEvaluator bsonTemplateEvaluator;
+        invariant(bsonTemplateEvaluator.setId(_id) == BsonTemplateEvaluator::StatusSuccess);
 
         while ( !shouldStop() ) {
             BSONObjIterator i( _config->ops );
@@ -538,23 +539,24 @@ namespace mongo {
                         {
                             BenchRunEventTrace _bret(&_stats.insertCounter);
 
+                            BSONObjBuilder builder;
+                            builder.append("insert",
+                                nsToCollectionSubstring(ns));
+                            BSONArrayBuilder docBuilder(
+                                builder.subarrayStart("documents"));
+                            docBuilder.append(fixQuery(e["doc"].Obj(),
+                                bsonTemplateEvaluator));
+                            docBuilder.done();
+                            BSONObj insertObj = builder.obj();
+
                             if (useWriteCmd) {
                                 // TODO: Replace after SERVER-11774.
-                                BSONObjBuilder builder;
-                                builder.append("insert",
-                                    nsToCollectionSubstring(ns));
-                                BSONArrayBuilder docBuilder(
-                                    builder.subarrayStart("documents"));
-                                docBuilder.append(fixQuery(e["doc"].Obj(),
-                                    bsonTemplateEvaluator));
-                                docBuilder.done();
                                 conn->runCommand(
                                     nsToDatabaseSubstring(ns).toString(),
-                                    builder.obj(), result);
+                                    insertObj, result);
                             }
                             else {
-                                conn->insert(ns, fixQuery(e["doc"].Obj(),
-                                    bsonTemplateEvaluator));
+                                conn->insert(ns, insertObj);
                                 if (safe)
                                     result = conn->getLastErrorDetailed();
                             }
@@ -766,7 +768,7 @@ namespace mongo {
 
          // Start threads
          for ( unsigned i = 0; i < _config->parallel; i++ ) {
-             BenchRunWorker *worker = new BenchRunWorker(_config.get(), &_brState);
+             BenchRunWorker *worker = new BenchRunWorker(i, _config.get(), &_brState);
              worker->start();
              _workers.push_back(worker);
          }
diff --git a/src/mongo/scripting/bench.h b/src/mongo/scripting/bench.h
index b3a644c7765..098965de24e 100644
--- a/src/mongo/scripting/bench.h
+++ b/src/mongo/scripting/bench.h
@@ -305,8 +305,10 @@ namespace mongo {
          * Create a new worker, performing one thread's worth of the activity described in
          * "config", and part of the larger activity with state "brState".  Both "config"
          * and "brState" must exist for the life of this object.
+         *
+         * "id" is a positive integer which should uniquely identify the worker.
          */
-        BenchRunWorker(const BenchRunConfig *config, BenchRunState *brState);
+        BenchRunWorker(size_t id, const BenchRunConfig *config, BenchRunState *brState);
         ~BenchRunWorker();
 
         /**
@@ -332,6 +334,7 @@ namespace mongo {
         /// Predicate, used to decide whether or not it's time to terminate the worker.
         bool shouldStop() const;
 
+        size_t _id;
         const BenchRunConfig *_config;
         BenchRunState *_brState;
         BenchRunStats _stats;
diff --git a/src/mongo/scripting/bson_template_evaluator.cpp b/src/mongo/scripting/bson_template_evaluator.cpp
index 4d044cc0590..038ab4582fd 100644
--- a/src/mongo/scripting/bson_template_evaluator.cpp
+++ b/src/mongo/scripting/bson_template_evaluator.cpp
@@ -32,19 +32,30 @@
 #include <cstdlib>
 
 #include "mongo/util/map_util.h"
+#include "mongo/util/mongoutils/str.h"
 
 namespace mongo {
 
     void BsonTemplateEvaluator::initializeEvaluator() {
         addOperator("RAND_INT", &BsonTemplateEvaluator::evalRandInt);
+        addOperator("SEQ_INT", &BsonTemplateEvaluator::evalSeqInt);
         addOperator("RAND_STRING", &BsonTemplateEvaluator::evalRandString);
         addOperator("CONCAT", &BsonTemplateEvaluator::evalConcat);
+        addOperator("OID", &BsonTemplateEvaluator::evalObjId);
     }
 
-    BsonTemplateEvaluator::BsonTemplateEvaluator() {
+    BsonTemplateEvaluator::BsonTemplateEvaluator() : _id(0) {
         initializeEvaluator();
     }
 
+    BsonTemplateEvaluator::Status BsonTemplateEvaluator::setId(size_t id) {
+        if (id >= (1<<7))
+            // The id is too large--doesn't fit inside 7 bits.
+            return StatusOpEvaluationError;
+        this->_id = id;
+        return StatusSuccess;
+    }
+
     BsonTemplateEvaluator::~BsonTemplateEvaluator() { }
 
     void BsonTemplateEvaluator::addOperator(const std::string& name, const OperatorFn& op) {
@@ -94,7 +105,7 @@ namespace mongo {
 
     BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandInt(BsonTemplateEvaluator* btl,
                                                                      const char* fieldName,
-                                                                     const BSONObj in,
+                                                                     const BSONObj& in,
                                                                      BSONObjBuilder& out) {
         // in = { #RAND_INT: [10, 20] }
         BSONObj range = in.firstElement().embeddedObject();
@@ -114,9 +125,62 @@ namespace mongo {
         return StatusSuccess;
     }
 
+    BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalSeqInt(BsonTemplateEvaluator* btl,
+                                                                    const char* fieldName,
+                                                                    const BSONObj& in,
+                                                                    BSONObjBuilder& out) {
+        // in = { #SEQ_INT: { seq_id: 0, start: 10, step: -2 }
+        BSONObj spec = in.firstElement().embeddedObject();
+        if (spec.nFields() < 3)
+            return StatusOpEvaluationError;
+        if (spec["seq_id"].eoo() || !spec["seq_id"].isNumber())
+            return StatusOpEvaluationError;
+        if (spec["start"].eoo() || !spec["start"].isNumber())
+            return StatusOpEvaluationError;
+        if (spec["step"].eoo() || !spec["step"].isNumber())
+            return StatusOpEvaluationError;
+
+        // If we're here, then we have a well-formed SEQ_INT specification:
+        // seq_id, start, and step fields, which are all numbers.
+
+        int seq_id = spec["seq_id"].numberInt();
+        long long curr_seqval = spec["start"].numberInt();
+
+        // Handle the optional "unique" argument.
+        //
+        // If the test requires us to keep sequences unique between different
+        // worker threads, then put the id number of this evaluator in the
+        // high order byte.
+        if (!spec["unique"].eoo() && spec["unique"].trueValue()) {
+            long long workerid = btl->_id;
+            curr_seqval += (workerid << (64-8));
+        }
+
+        if (btl->_seqIdMap.end() != btl->_seqIdMap.find(seq_id)) {
+            // We already have a sequence value. Add 'step' to get the next value.
+            int step = spec["step"].numberInt();
+            curr_seqval = btl->_seqIdMap[seq_id] + step;
+        }
+
+        // Handle the optional "mod" argument. This should be done after
+        // handling all other options (currently just "unique").
+        if (!spec["mod"].eoo()) {
+            if (!spec["mod"].isNumber())
+                return StatusOpEvaluationError;
+            int modval = spec["mod"].numberInt();
+            curr_seqval = (curr_seqval % modval);
+        }
+
+        // Store the sequence value.
+        btl->_seqIdMap[seq_id] = curr_seqval;
+
+        out.append(fieldName, curr_seqval);
+        return StatusSuccess;
+    }
+
     BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalRandString(BsonTemplateEvaluator* btl,
                                                                         const char* fieldName,
-                                                                        const BSONObj in,
+                                                                        const BSONObj& in,
                                                                         BSONObjBuilder& out) {
         // in = { #RAND_STRING: [10] }
         BSONObj range = in.firstElement().embeddedObject();
@@ -145,7 +209,7 @@ namespace mongo {
 
     BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalConcat(BsonTemplateEvaluator* btl,
                                                                     const char* fieldName,
-                                                                    const BSONObj in,
+                                                                    const BSONObj& in,
                                                                     BSONObjBuilder& out) {
         // in = { #CONCAT: ["hello", " ", "world"] }
         BSONObjBuilder objectBuilder;
@@ -166,4 +230,16 @@ namespace mongo {
         return StatusSuccess;
     }
 
+    BsonTemplateEvaluator::Status BsonTemplateEvaluator::evalObjId(BsonTemplateEvaluator* btl,
+                                                                   const char* fieldName,
+                                                                   const BSONObj& in,
+                                                                   BSONObjBuilder& out) {
+        // in = { #OID: 1 }
+        if (!mongoutils::str::equals(fieldName, "_id"))
+            // Error: must be generating a value for the _id field.
+            return StatusOpEvaluationError;
+        out.genOID();
+        return StatusSuccess;
+    }
+
 } // end namespace mongo
diff --git a/src/mongo/scripting/bson_template_evaluator.h b/src/mongo/scripting/bson_template_evaluator.h
index ada4547f693..b6332440b82 100644
--- a/src/mongo/scripting/bson_template_evaluator.h
+++ b/src/mongo/scripting/bson_template_evaluator.h
@@ -17,7 +17,7 @@
 /*
  * This library supports a templating language that helps in generating BSON documents from a
  * template. The language supports the following template:
- * #RAND_INT, #RAND_STRING and #CONCAT.
+ * #RAND_INT, #SEQ_INT, #RAND_STRING, #CONCAT, and #OID.
  *
  * The language will help in quickly expressing richer documents  for use in benchRun.
  * Ex. : { key : { #RAND_INT: [10, 20] } } or  { key : { #CONCAT: ["hello", " ", "world"] } }
@@ -88,6 +88,11 @@ namespace mongo {
         BsonTemplateEvaluator();
         ~BsonTemplateEvaluator();
 
+        /**
+         * Set an identifying number for this template evaluator.
+         */
+        Status setId(size_t id);
+
         /*
          * "Add a new operator, "name" with behavior "op" to this evaluator.
          */
@@ -117,6 +122,15 @@ namespace mongo {
         // evaluates a BSON element. This is internally called by the top level evaluate method.
         Status _evalElem(BSONElement in, BSONObjBuilder& out);
 
+        // An identifier for this template evaluator instance, which distinguishes it
+        // from other evaluators. Useful for threaded benchruns which, say, want to insert
+        // similar sequences of values without colliding with one another.
+        unsigned char _id;
+
+        // Keeps state for each SEQ_INT expansion being evaluated by this bson template evaluator
+        // instance. Maps from the seq_id of the sequence to its current value.
+        std::map< int, long long > _seqIdMap;
+
         /*
          * Operator method to support #RAND_INT :  { key : { #RAND_INT: [10, 20] } }
          * The array arguments to #RAND_INT are the min and mix range between which a random number
@@ -128,14 +142,38 @@ namespace mongo {
          * choose a random number between 10 and 20 and then multiple the chosen value with 4.
          */
         static Status evalRandInt(BsonTemplateEvaluator* btl, const char* fieldName,
-                                  const BSONObj in, BSONObjBuilder& out);
+                                  const BSONObj& in, BSONObjBuilder& out);
+        /*
+         * Operator method to support #SEQ_INT :
+         *    { key : { #SEQ_INT: { seq_id: 0, start: 100, step: -2, unique: true } } }
+         *
+         * Used to generate arithmetic sequences of integers in each successive template
+         * evaluation. The 'seq_id' identifies this expansion so that the same document can
+         * have multiple expansions. The sequence will begin with the value 'start' and
+         * then increment by 'step' for each successive expansion.
+         *
+         * If 'unique: true' is specified, then the sequences are adjusted according to the
+         * _id of this template evaluator so that various worker threads don't overrun each
+         * other's ranges.
+         *
+         * If using with multiple threads, each thread should have its own BsonTemplateEvaluator
+         * instance so that the sequence state remains separate.
+         *
+         * If 'mod: <num>' is specified, then modulo <num> is applied to each value of the
+         * arithmetic sequence.
+         *
+         * Ex. { a: { #SEQ_INT: { seq_id: 0, start: 4, step: 3 } } } will generate
+         *    { a: 4 }, { a: 7 }, { a: 10 }, etc.
+         */
+        static Status evalSeqInt(BsonTemplateEvaluator* btl, const char* fieldName,
+                                 const BSONObj& in, BSONObjBuilder& out);
         /*
          * Operator method to support #RAND_STRING : { key : { #RAND_STRING: [12] } }
          * The array argument to RAND_STRING is the length of the string that is desired.
          * This will evaluate to something like { key : "randomstring" }
          */
         static Status evalRandString(BsonTemplateEvaluator* btl, const char* fieldName,
-                                     const BSONObj in, BSONObjBuilder& out);
+                                     const BSONObj& in, BSONObjBuilder& out);
         /*
          * Operator method to support #CONCAT : { key : { #CONCAT: ["hello", " ", "world", 2012] } }
          * The array argument to CONCAT are the strings to be concatenated. If the argument is not
@@ -143,7 +181,14 @@ namespace mongo {
          * This will evaluate to { key : "hello world2012" }
          */
         static Status evalConcat(BsonTemplateEvaluator* btl, const char* fieldName,
-                                 const BSONObj in, BSONObjBuilder& out);
+                                 const BSONObj& in, BSONObjBuilder& out);
+        /*
+         * Operator method to support #OID : { _id : { #OID: 1 } }
+         * The 'key' field is required to be _id, but the argument to OID does not matter,
+         * and will be neither examined nor validated.
+         */
+        static Status evalObjId(BsonTemplateEvaluator* btl, const char* fieldName,
+                                const BSONObj& in, BSONObjBuilder& out);
 
     };
 
diff --git a/src/mongo/scripting/bson_template_evaluator_test.cpp b/src/mongo/scripting/bson_template_evaluator_test.cpp
index cbd3dd07c1a..3248881a61b 100644
--- a/src/mongo/scripting/bson_template_evaluator_test.cpp
+++ b/src/mongo/scripting/bson_template_evaluator_test.cpp
@@ -173,6 +173,142 @@ namespace mongo {
             ASSERT_LESS_THAN_OR_EQUALS(obj12.firstElement().numberInt(), 16);
         }
 
+        TEST(BSONTemplateEvaluatorTest, SEQ_INT) {
+
+            boost::scoped_ptr<BsonTemplateEvaluator> t(new BsonTemplateEvaluator());
+            BSONObj seqObj;
+            BSONObj expectedObj;
+
+            // Error if missing 'step'.
+            BSONObjBuilder builder1;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "start" << 0 ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder1) );
+
+            // Error if missing 'start'.
+            BSONObjBuilder builder2;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "step" << 1 ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder2) );
+
+            // Error if missing 'seq_iq'.
+            BSONObjBuilder builder3;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "start" << 0 << "step" << 1 ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder3) );
+
+            // Error if 'step' is not a number.
+            BSONObjBuilder builder4;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "start" << 0 << "step" << "foo" ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder4) );
+
+            // Error if 'start' is not a number.
+            BSONObjBuilder builder5;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "start" << true << "step" << 1 ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder5) );
+
+            // Error if 'seq_id' is not a number.
+            BSONObjBuilder builder6;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << BSON("foo" << 1) << "start" << 0 << "step" << 1 ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder6) );
+
+            // Error if 'mod' is not a number.
+            BSONObjBuilder builder7;
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "start" << 0 << "step" << 1 << "mod" << "foo" ));
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("seqField" << seqObj), builder7) );
+
+            // Test an increasing sequence: -4, -2, 0, 2, ...
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 0 << "start" << -4 << "step" << 2 ));
+            for (int i = -4; i <= 2; i += 2) {
+                BSONObjBuilder builder8;
+                ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                               t->evaluate(BSON("seqField" << seqObj), builder8) );
+                expectedObj = BSON("seqField" << i);
+                ASSERT_EQUALS(0, expectedObj.woCompare(builder8.obj()));
+            }
+
+            // Test a decreasing sequence: 5, 0, -5, -10
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 1 << "start" << 5 << "step" << -5 ));
+            for (int i = 5; i >= -10; i -= 5) {
+                BSONObjBuilder builder9;
+                ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                               t->evaluate(BSON("seqField" << seqObj), builder9) );
+                expectedObj = BSON("seqField" << i);
+                ASSERT_EQUALS(0, expectedObj.woCompare(builder9.obj()));
+            }
+
+            // Test multiple sequences in the same document. In order for this to
+            // work the two sequences must have different sequence IDs.
+            //
+            // seq_id 2: 0, 1, 2, 3, ...
+            // seq_id 3: 0, -1, -2, -3, ...
+            BSONObj seqObj1 = BSON( "#SEQ_INT"
+                              << BSON( "seq_id" << 2 << "start" << 0 << "step" << 1 ));
+            BSONObj seqObj2 = BSON( "#SEQ_INT"
+                              << BSON( "seq_id" << 3 << "start" << 0 << "step" << -1 ));
+            BSONObj seqObjFull = BSON( "seqField1" << seqObj1 << "seqField2" << seqObj2 );
+            for (int i = 0; i <= 3; i++) {
+                BSONObjBuilder builder10;
+                ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                               t->evaluate(seqObjFull, builder10));
+                expectedObj = BSON("seqField1" << i << "seqField2" << -i);
+                ASSERT_EQUALS(0, expectedObj.woCompare(builder10.obj()));
+            }
+
+            // Test that the 'unique: true' option correctly puts the ID of the
+            // bson template evaluator into the high order byte of a 64 bit integer.
+            t->setId(9);
+            seqObj1 = BSON( "#SEQ_INT"
+                      << BSON( "seq_id" << 4 << "start" << 8 << "step" << 1 ));
+            seqObj2 = BSON( "#SEQ_INT"
+                      << BSON( "seq_id" << 5 << "start" << 8 << "step" << 1
+                      << "unique" << true ));
+
+            // Without 'unique: true'.
+            BSONObjBuilder builder11;
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                           t->evaluate(BSON("seqField" << seqObj1), builder11) );
+            expectedObj = BSON("seqField" << 8);
+            ASSERT_EQUALS(0, expectedObj.woCompare(builder11.obj()));
+
+            // With 'unique: true'.
+            BSONObjBuilder builder12;
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                           t->evaluate(BSON("seqField" << seqObj2), builder12) );
+            // The template evaluator id of 9 goes in the high-order byte.
+            long long expectedSeqNum = 0x0900000000000008;
+            expectedObj = BSON("seqField" << expectedSeqNum);
+            ASSERT_EQUALS(0, expectedObj.woCompare(builder12.obj()));
+
+            // Test a sequence using "mod": 0, 1, 2, 0, 1
+            seqObj = BSON( "#SEQ_INT"
+                     << BSON( "seq_id" << 6 << "start" << 0 << "step" << 1 << "mod" << 3 ));
+            for (int i = 0; i <= 5; i++) {
+                BSONObjBuilder builder13;
+                ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                               t->evaluate(BSON("seqField" << seqObj), builder13) );
+                expectedObj = BSON("seqField" << (i%3));
+                ASSERT_EQUALS(0, expectedObj.woCompare(builder13.obj()));
+            }
+
+            // Test that you can't set an id if it is more than 7 bits wide.
+            ASSERT_EQUALS(BsonTemplateEvaluator::StatusSuccess, t->setId(127));
+            ASSERT_EQUALS(BsonTemplateEvaluator::StatusOpEvaluationError, t->setId(128));
+        }
+
         TEST(BSONTemplateEvaluatorTest, RAND_STRING) {
 
             BsonTemplateEvaluator *t = new BsonTemplateEvaluator();
@@ -333,6 +469,23 @@ namespace mongo {
             ASSERT_EQUALS(obj5.equal(expectedObj), true);
         }
 
+        TEST(BSONTemplateEvaluatorTest, OID) {
+
+            boost::scoped_ptr<BsonTemplateEvaluator> t(new BsonTemplateEvaluator());
+            BSONObj oidObj = BSON( "#OID" << 1 );
+
+            // Error: field must be "_id"
+            BSONObjBuilder builder1;
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusOpEvaluationError,
+                           t->evaluate(BSON("notIdField" << oidObj), builder1) );
+
+            // Success.
+            BSONObjBuilder builder2;
+            ASSERT_EQUALS( BsonTemplateEvaluator::StatusSuccess,
+                           t->evaluate(BSON("_id" << oidObj), builder2) );
+
+        }
+
         TEST(BSONTemplateEvaluatorTest, COMBINED_OPERATORS) {
 
             BsonTemplateEvaluator *t = new BsonTemplateEvaluator();