Add cql connect/read options

yugabyte · Apr 24, 2020 · f9866b6 · f9866b6
1 parent 55b8df9
commit f9866b6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 3 deletions.
diff --git a/src/main/java/com/yugabyte/sample/apps/AppBase.java b/src/main/java/com/yugabyte/sample/apps/AppBase.java
@@ -43,6 +43,7 @@
 import com.datastax.driver.core.SimpleStatement;
 import com.datastax.driver.core.SSLOptions;
 import com.datastax.driver.core.JdkSSLOptions;
+import com.datastax.driver.core.SocketOptions;
 import com.datastax.driver.core.exceptions.NoHostAvailableException;
 import com.datastax.driver.core.policies.DCAwareRoundRobinPolicy;
 import com.datastax.driver.core.policies.DefaultRetryPolicy;
@@ -242,6 +243,14 @@ protected synchronized void createCassandraClient(List<ContactPoint> contactPoin
             .withSSL(createSSLHandler(appConfig.sslCert));
       }
       Integer port = null;
+      SocketOptions socketOptions = new SocketOptions();
+      if (appConfig.cqlConnectTimeoutMs > 0) {
+        socketOptions.setConnectTimeoutMillis(appConfig.cqlConnectTimeoutMs);
+      }
+      if (appConfig.cqlReadTimeoutMs > 0) {
+        socketOptions.setReadTimeoutMillis(appConfig.cqlReadTimeoutMs);
+      }
+      builder.withSocketOptions(socketOptions);
       for (ContactPoint cp : contactPoints) {
         if (port == null) {
           port = cp.getPort();
@@ -375,13 +384,12 @@ protected byte[] getRandomValue(Key key) {
     getRandomValue(key, buffer);
     return buffer;
   }
-
 
   protected byte[] getRandomValue(Key key, byte[] outBuffer) {
     getRandomValue(key, outBuffer.length, outBuffer);
     return outBuffer;
   }
-  
+
   protected void getRandomValue(Key key, int valueSize, byte[] outBuffer) {
 	  final byte[] keyValueBytes = key.getValueStr().getBytes();
 	  getRandomValue(keyValueBytes, valueSize, outBuffer);

diff --git a/src/main/java/com/yugabyte/sample/apps/AppConfig.java b/src/main/java/com/yugabyte/sample/apps/AppConfig.java
@@ -59,6 +59,9 @@ public static enum Type {
   public int sleepTime = 0;
   public int jedisSocketTimeout = 61000;
 
+  public int cqlConnectTimeoutMs = 0;
+  public int cqlReadTimeoutMs = 0;
+
   // Use ASCII strings as values.
   public boolean restrictValuesToAscii;
 

diff --git a/src/main/java/com/yugabyte/sample/common/CmdLineOpts.java b/src/main/java/com/yugabyte/sample/common/CmdLineOpts.java
@@ -505,6 +505,14 @@ private void initializeNumKeys(CommandLine cmd) {
       AppBase.appConfig.jedisSocketTimeout =
           Integer.parseInt(cmd.getOptionValue("socket_timeout"));
     }
+    if (cmd.hasOption("cql_connect_timeout_ms")) {
+      AppBase.appConfig.cqlConnectTimeoutMs =
+          Integer.parseInt(cmd.getOptionValue("cql_connect_timeout_ms"));
+    }
+    if (cmd.hasOption("cql_read_timeout_ms")) {
+      AppBase.appConfig.cqlReadTimeoutMs =
+          Integer.parseInt(cmd.getOptionValue("cql_read_timeout_ms"));
+    }
     if (cmd.hasOption("use_ascii_values")) {
       AppBase.appConfig.restrictValuesToAscii = true;
     }
@@ -600,6 +608,8 @@ public static CmdLineOpts createFromArgs(String[] args) throws Exception {
         "How long (in ms) to sleep between multiple pipeline batches.");
     options.addOption("socket_timeout", true,
                       "How long (in ms) to wait for a response from jedis.");
+    options.addOption("cql_connect_timeout_ms", true, "Connection timeout for cql in millisecs");
+    options.addOption("cql_read_timeout_ms", true, "Read timeout for cql in millisecs");
     options.addOption("value_size", true, "Size in bytes of the value. " +
         "The bytes are random. Value size should be more than 5 (9) bytes for binary (ascii) " +
         "values in order to have checksum for read verification. First byte is used as a " +
@@ -625,7 +635,7 @@ public static CmdLineOpts createFromArgs(String[] args) throws Exception {
             "If this option is set, the --username option is required.");
     options.addOption("concurrent_clients", true,
         "The number of client connections to establish to each host in the YugaByte DB cluster.");
-    options.addOption("ssl_cert", true, 
+    options.addOption("ssl_cert", true,
       "Use an SSL connection while connecting to YugaByte.");
     options.addOption("batch_size", true,
                       "Number of keys to write in a batch (for apps that support batching).");