SERVER-20297 Expose cluster nodes to workload states in the concurrency suite

Closes #1036

Signed-off-by: Kamran Khan <kamran.khan@mongodb.com>

5 files changed, 144 insertions, 11 deletions

diff --git a/jstests/concurrency/fsm_libs/cluster.js b/jstests/concurrency/fsm_libs/cluster.js
index 99236dc29d6..86265ca3fbf 100644
--- a/jstests/concurrency/fsm_libs/cluster.js
+++ b/jstests/concurrency/fsm_libs/cluster.js
@@ -65,6 +65,8 @@ var Cluster = function(options) {
 
     var conn;
 
+    var st;
+
     var initialized = false;
 
     var _conns = {
@@ -108,8 +110,7 @@ var Cluster = function(options) {
                 };
             }
 
-            var st = new ShardingTest(shardConfig);
-            st.stopBalancer();
+            st = new ShardingTest(shardConfig);
 
             conn = st.s; // mongos
 
@@ -147,7 +148,6 @@ var Cluster = function(options) {
                 ++i;
                 mongod = st['d' + i];
             }
-
         } else if (options.replication) {
             // TODO: allow 'options' to specify the number of nodes
             var replSetConfig = {
@@ -197,7 +197,7 @@ var Cluster = function(options) {
         }
 
         initialized = true;
-        
+
         this.executeOnMongodNodes(options.setupFunctions.mongod);
         this.executeOnMongosNodes(options.setupFunctions.mongos);
     };
@@ -217,12 +217,12 @@ var Cluster = function(options) {
         }
         if (!fn || typeof(fn) !== 'function' || fn.length !== 1) {
             throw new Error('mongod function must be a function that takes a db as an argument');
-        } 
+        }
         _conns.mongod.forEach(function(mongodConn) {
             fn(mongodConn.getDB('admin'));
         });
     };
-        
+
     this.executeOnMongosNodes = function executeOnMongosNodes(fn) {
         if (!initialized) {
             throw new Error('cluster must be initialized before functions can be executed ' +
@@ -230,12 +230,12 @@ var Cluster = function(options) {
         }
         if (!fn || typeof(fn) !== 'function' || fn.length !== 1) {
             throw new Error('mongos function must be a function that takes a db as an argument');
-        } 
+        }
         _conns.mongos.forEach(function(mongosConn) {
             fn(mongosConn.getDB('admin'));
         });
     };
-        
+
     this.teardown = function teardown() { };
 
     this.getDB = function getDB(dbName) {
@@ -271,6 +271,91 @@ var Cluster = function(options) {
         throw new Error('cluster has not been initialized yet');
     };
 
+    // Provide a serializable form of the cluster for use in workload states. This
+    // method is required because we don't currently support the serialization of Mongo
+    // connection objects.
+    //
+    // Serialized format:
+    // {
+    //      mongos: [
+    //          "localhost:30998",
+    //          "localhost:30999"
+    //      ],
+    //      config: [
+    //          "localhost:29000",
+    //          "localhost:29001",
+    //          "localhost:29002"
+    //      ],
+    //      shards: {
+    //          "test-rs0": [
+    //              "localhost:20006",
+    //              "localhost:20007",
+    //              "localhost:20008"
+    //          ],
+    //          "test-rs1": [
+    //              "localhost:20009",
+    //              "localhost:20010",
+    //              "localhost:20011"
+    //          ]
+    //      }
+    // }
+    this.getSerializedCluster = function getSerializedCluster() {
+        // TODO: Add support for non-sharded clusters.
+        if (!this.isSharded()) {
+            return '';
+        }
+
+        var cluster = {
+            mongos: [],
+            config: [],
+            shards: {}
+        };
+
+        var i = 0;
+        var mongos = st.s0;
+        while (mongos) {
+            cluster.mongos.push(mongos.name);
+            ++i;
+            mongos = st['s' + i];
+        }
+
+        i = 0;
+        var config = st.c0;
+        while (config) {
+            cluster.config.push(config.name);
+            ++i;
+            config = st['c' + i];
+        }
+
+        i = 0;
+        var shard = st.shard0;
+        while (shard) {
+            if (shard.name.includes('/')) {
+                // If the shard is a replica set, the format of st.shard0.name in ShardingTest is
+                // "test-rs0/localhost:20006,localhost:20007,localhost:20008".
+                var [setName, shards] = shard.name.split('/');
+                cluster.shards[setName] = shards.split(',');
+            } else {
+                // If the shard is a standalone mongod, the format of st.shard0.name in ShardingTest
+                // is "localhost:20006".
+                cluster.shards[i] = [shard.name];
+            }
+            ++i;
+            shard = st['shard' + i];
+        }
+        return cluster;
+    }
+
+    this.startBalancer = function startBalancer() {
+        assert(this.isSharded(), 'cluster is not sharded');
+        st.startBalancer();
+    };
+
+    this.stopBalancer = function stopBalancer() {
+        assert(this.isSharded(), 'cluster is not sharded');
+        st.stopBalancer();
+    };
+
     this.awaitReplication = function awaitReplication() {
         if (this.isReplication()) {
             var wc = {
diff --git a/jstests/concurrency/fsm_libs/fsm.js b/jstests/concurrency/fsm_libs/fsm.js
index 2bb067df199..7ba2544f3fc 100644
--- a/jstests/concurrency/fsm_libs/fsm.js
+++ b/jstests/concurrency/fsm_libs/fsm.js
@@ -4,6 +4,8 @@ var fsm = (function() {
     // args.data = 'this' object of the state functions
     // args.db = database object
     // args.collName = collection name
+    // args.cluster = connection strings for all cluster nodes (see fsm_libs/cluster.js for format)
+    // args.passConnectionCache = boolean, whether to pass a connection cache to the workload states
     // args.startState = name of initial state function
     // args.states = state functions of the form
     //               { stateName: function(db, collName) { ... } }
@@ -13,13 +15,44 @@ var fsm = (function() {
     // args.iterations = number of iterations to run the FSM for
     function runFSM(args) {
         var currentState = args.startState;
+
+        // We build a cache of connections that can be used in workload states. This cache
+        // allows state functions to access arbitrary cluster nodes for verification checks.
+        // See fsm_libs/cluster.js for the format of args.cluster.
+        var connCache;
+        if (args.passConnectionCache) {
+            connCache = {
+                mongos: [],
+                config: [],
+                shards: {}
+            };
+            connCache.mongos = args.cluster.mongos.map(connStr => new Mongo(connStr));
+            connCache.config = args.cluster.config.map(connStr => new Mongo(connStr));
+
+            var shardNames = Object.keys(args.cluster.shards);
+
+
+            shardNames.forEach(name =>
+                connCache.shards[name] = args.cluster.shards[name].map(connStr =>
+                    new Mongo(connStr)));
+        }
+
         for (var i = 0; i < args.iterations; ++i) {
             var fn = args.states[currentState];
             assert.eq('function', typeof fn, 'states.' + currentState + ' is not a function');
-            fn.call(args.data, args.db, args.collName);
+            fn.call(args.data, args.db, args.collName, connCache);
             var nextState = getWeightedRandomChoice(args.transitions[currentState], Random.rand());
             currentState = nextState;
         }
+
+        if (args.passConnectionCache) {
+            connCache.mongos.forEach(conn => conn = null);
+            connCache.config.forEach(conn => conn = null);
+
+            var shardNames = Object.keys(connCache.shards);
+            shardNames.forEach(name => connCache.shards[name].forEach(conn => conn = null));
+            gc();
+        }
     }
 
     // doc = document of the form
diff --git a/jstests/concurrency/fsm_libs/parse_config.js b/jstests/concurrency/fsm_libs/parse_config.js
index fa482e83151..d1aef30c545 100644
--- a/jstests/concurrency/fsm_libs/parse_config.js
+++ b/jstests/concurrency/fsm_libs/parse_config.js
@@ -9,6 +9,7 @@ function parseConfig(config) {
     var allowedKeys = [
         'data',
         'iterations',
+        'passConnectionCache',
         'setup',
         'startState',
         'states',
@@ -39,8 +40,15 @@ function parseConfig(config) {
     Object.keys(config.states).forEach(function(k) {
         assert.eq('function', typeof config.states[k],
                    'config.states.' + k + ' is not a function');
-        assert.eq(2, config.states[k].length,
-                  'state functions should accept 2 parameters: db and collName');
+        if (config.passConnectionCache) {
+             assert.eq(3, config.states[k].length,
+                      'if passConnectionCache is true, state functions should ' +
+                      'accept 3 parameters: db, collName, and connCache');
+        } else {
+            assert.eq(2, config.states[k].length,
+                      'if passConnectionCache is false, state functions should ' +
+                      'accept 2 parameters: db and collName');
+        }
     });
 
     // assert all states mentioned in config.transitions are present in config.states
@@ -72,5 +80,8 @@ function parseConfig(config) {
     assert.eq('object', typeof config.data);
     // TODO: assert that 'tid' is not a key of 'config.data'
 
+    config.passConnectionCache = config.passConnectionCache || false;
+    assert.eq('boolean', typeof config.passConnectionCache);
+
     return config;
 }
diff --git a/jstests/concurrency/fsm_libs/thread_mgr.js b/jstests/concurrency/fsm_libs/thread_mgr.js
index 6073d82dd2a..86737a1a11d 100644
--- a/jstests/concurrency/fsm_libs/thread_mgr.js
+++ b/jstests/concurrency/fsm_libs/thread_mgr.js
@@ -116,6 +116,7 @@ var ThreadManager = function(clusterOptions, executionMode) {
                     latch: latch,
                     dbName: _context[workload].dbName,
                     collName: _context[workload].collName,
+                    cluster: cluster.getSerializedCluster(),
                     clusterOptions: clusterOptions,
                     seed: Random.randInt(1e13), // contains range of Date.getTime()
                     globalAssertLevel: globalAssertLevel
diff --git a/jstests/concurrency/fsm_libs/worker_thread.js b/jstests/concurrency/fsm_libs/worker_thread.js
index c2e4cb87509..ec6118a6f08 100644
--- a/jstests/concurrency/fsm_libs/worker_thread.js
+++ b/jstests/concurrency/fsm_libs/worker_thread.js
@@ -13,6 +13,7 @@ var workerThread = (function() {
     // args.latch = CountDownLatch instance for starting all threads
     // args.dbName = the database name
     // args.collName = the collection name
+    // args.cluster = connection strings for all cluster nodes (see cluster.js for format)
     // args.clusterOptions = the configuration of the cluster
     // args.seed = seed for the random number generator
     // args.globalAssertLevel = the global assertion level to use
@@ -61,6 +62,8 @@ var workerThread = (function() {
                     data: data,
                     db: myDB,
                     collName: args.collName,
+                    cluster: args.cluster,
+                    passConnectionCache: config.passConnectionCache,
                     startState: config.startState,
                     states: config.states,
                     transitions: config.transitions,