Allow for multiple Cassandra/MQTT hosts to be specified, picking one randomly to connect to

tools/dcdbslurmjob/dcdbslurmjob.cpp

@@ -62,8 +62,8 @@ void usage() {
     std::cout << std::endl;
 
     std::cout << "Options:" << std::endl;
-    std::cout << "  -b<host>      MQTT broker                    [default: 127.0.0.1:1883]" << std::endl;
-    std::cout << "  -c<host>      Cassandra host                 [default: 127.0.0.1:9042]" << std::endl;
+    std::cout << "  -b<hosts>     List of MQTT brokers           [default: localhost:1883]" << std::endl;
+    std::cout << "  -c<hosts>     List of Cassandra hosts        [default: none]" << std::endl;
     std::cout << "  -u<username>  Cassandra username             [default: none]" << std::endl;
     std::cout << "  -p<password>  Cassandra password             [default: none]" << std::endl;
     std::cout << "  -t<timestamp> Timestamp value                [default: now]" << std::endl;
@@ -136,6 +136,23 @@ void convertNodeList(DCDB::NodeList& nl, std::string substitution) {
     }
 }
 
+void splitHostList(const std::string& str, std::vector<std::string>& hl, char delim = ',')
+{
+    hl.clear();
+    std::stringstream ss(str);
+    std::string token;
+    while (std::getline(ss, token, delim)) {
+	hl.push_back(token);
+    }
+}
+
+void pickRandomHost(const std::vector<std::string>& hl, std::string& host, int& port) {
+    srand (time(NULL));
+    int n = rand() % hl.size();
+    host = parseNetworkHost(hl[n]);
+    port = atoi(parseNetworkPort(hl[n]).c_str());
+}
+
 /**
  * Retrieves Slurm job data from environment variables and sends it to either a
  * CollectAgent or a Cassandra database. Job data can also be passed as command
@@ -149,8 +166,9 @@ int main(int argc, char** argv) {
     DCDB::JobDataStore *myJobDataStore = nullptr;
     struct mosquitto *  _mosq = nullptr;
 
-    std::string host = "127.0.0.1", cassandraPort = "9042", cassandraUser = "", cassandraPassword = "";
-    int         brokerPort = 1883;
+    std::vector<std::string> hostList;
+    std::string host = "", cassandraUser = "", cassandraPassword = "";
+    int         port;
     std::string nodelist="", jobId="", userId="";
     std::string substitution="";
     uint64_t ts=0;
@@ -183,21 +201,12 @@ int main(int argc, char** argv) {
         switch(ret) {
 		case 'b': {
 			cassandra = false;
-			host = parseNetworkHost(optarg);
-			std::string port = parseNetworkPort(optarg);
-			if (port != "") {
-				brokerPort = std::stoi(port);
-			} else {
-				brokerPort = 1883;
-			}
+			splitHostList(optarg, hostList);
 			break;
 		}
 		case 'c':
 			cassandra = true;
-			host = parseNetworkHost(optarg);
-			cassandraPort = parseNetworkPort(optarg);
-			if (cassandraPort == "")
-				cassandraPort = std::string("9042");
+			splitHostList(optarg, hostList);
 			break;
 		case 'u':
 			cassandra = true;
@@ -238,15 +247,24 @@ int main(int argc, char** argv) {
                 return 1;
         }
     }
+    
+    if (hostList.size() == 0) {
+	hostList.push_back("localhost");
+    }
 
     if (cassandra) {
 	    //Allocate and initialize connection to Cassandra.
-	    dcdbConn = new DCDB::Connection(host, atoi(cassandraPort.c_str()), cassandraUser, cassandraPassword);
-
+	    pickRandomHost(hostList, host, port);
+	    if (port == 0) {
+		port = 9042;
+	    }
+	
+	    dcdbConn = new DCDB::Connection(host, port, cassandraUser, cassandraPassword);
 	    if (!dcdbConn->connect()) {
-		    std::cerr << "Cannot connect to Cassandra!" << std::endl;
+		    std::cerr << "Cannot connect to Cassandra server " << host << ":" << port << std::endl;
 		    return 1;
 	    }
+	    std::cout << "Connected to Cassandra server " << host << ":" << port << std::endl;
 	    myJobDataStore = new DCDB::JobDataStore(dcdbConn);
     } else {
 	    //Initialize Mosquitto library and connect to broker
@@ -264,10 +282,16 @@ int main(int argc, char** argv) {
 		    return 1;
 	    }
 
-	    if (mosquitto_connect(_mosq, host.c_str(), brokerPort, 1000) != MOSQ_ERR_SUCCESS) {
-		    std::cerr << "Could not connect to MQTT broker " << host << ":" << std::to_string(brokerPort) << std::endl;
+	    pickRandomHost(hostList, host, port);
+	    if (port == 0) {
+		port = 1883;
+	    }
+
+	    if (mosquitto_connect(_mosq, host.c_str(), port, 1000) != MOSQ_ERR_SUCCESS) {
+		    std::cerr << "Could not connect to MQTT broker " << host << ":" << port << std::endl;
 		    return 1;
 	    }
+	    std::cout << "Connected to MQTT broker " << host << ":" << port << std::endl;
     }
 
     //collect job data