Merge pull request #3052 from HubSpot/configurable_default_timeout

makes the various timeouts configurable on the connection -- connect,…
vitessio · Sep 3, 2017 · 76f5575 · 76f5575
2 parents fc4674f + cf529be
commit 76f5575
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 42 deletions.
diff --git a/java/client/src/main/java/io/vitess/client/cursor/FieldMap.java b/java/client/src/main/java/io/vitess/client/cursor/FieldMap.java
@@ -19,13 +19,17 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
-import com.google.common.collect.ImmutableList;
-import io.vitess.proto.Query.Field;
 import java.util.List;
 import java.util.Map;
+
 import javax.annotation.Nullable;
+
 import org.apache.commons.collections4.map.CaseInsensitiveMap;
 
+import com.google.common.collect.ImmutableList;
+
+import io.vitess.proto.Query.Field;
+
 /**
  * A wrapper for {@code List<Field>} that also facilitates lookup by field name.
  *

diff --git a/java/jdbc/src/main/java/io/vitess/jdbc/ConnectionProperties.java b/java/jdbc/src/main/java/io/vitess/jdbc/ConnectionProperties.java
@@ -16,11 +16,6 @@
 
 package io.vitess.jdbc;
 
-import io.vitess.proto.Query;
-import io.vitess.proto.Topodata;
-import io.vitess.util.Constants;
-import io.vitess.util.StringUtils;
-
 import java.io.UnsupportedEncodingException;
 import java.lang.reflect.Field;
 import java.sql.DriverPropertyInfo;
@@ -30,6 +25,11 @@
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 
+import io.vitess.proto.Query;
+import io.vitess.proto.Topodata;
+import io.vitess.util.Constants;
+import io.vitess.util.StringUtils;
+
 public class ConnectionProperties {
 
     private static final ArrayList<java.lang.reflect.Field> PROPERTY_LIST = new ArrayList<>();
@@ -207,6 +207,12 @@ public class ConnectionProperties {
         "Should the driver treat java.util.Date as a TIMESTAMP for the purposes of PreparedStatement.setObject()",
         true);
 
+    private LongConnectionProperty timeout = new LongConnectionProperty(
+        "timeout",
+        "The default timeout, in millis, to use for queries, connections, and transaction commit/rollback. Query timeout can be overridden by explicitly calling setQueryTimeout",
+        Constants.DEFAULT_TIMEOUT
+    );
+
     // Caching of some hot properties to avoid casting over and over
     private Topodata.TabletType tabletTypeCache;
     private Query.ExecuteOptions.IncludedFields includedFieldsCache;
@@ -462,6 +468,14 @@ public void setTreatUtilDateAsTimestamp(boolean treatUtilDateAsTimestamp) {
         this.treatUtilDateAsTimestamp.setValue(treatUtilDateAsTimestamp);
     }
 
+    public long getTimeout() {
+        return timeout.getValueAsLong();
+    }
+
+    public void setTimeout(long timeout) {
+        this.timeout.setValue(timeout);
+    }
+
     public String getTarget() {
         if(!StringUtils.isNullOrEmptyWithoutWS(target.getValueAsString())) {
             return target.getValueAsString();

diff --git a/java/jdbc/src/main/java/io/vitess/jdbc/VitessConnection.java b/java/jdbc/src/main/java/io/vitess/jdbc/VitessConnection.java
@@ -16,13 +16,6 @@
 
 package io.vitess.jdbc;
 
-import io.vitess.client.Context;
-import io.vitess.client.VTGateConn;
-import io.vitess.client.VTGateTx;
-import io.vitess.util.CommonUtils;
-import io.vitess.util.Constants;
-import io.vitess.util.MysqlDefs;
-
 import java.sql.Array;
 import java.sql.Blob;
 import java.sql.CallableStatement;
@@ -50,6 +43,13 @@
 import java.util.concurrent.Executor;
 import java.util.logging.Logger;
 
+import io.vitess.client.Context;
+import io.vitess.client.VTGateConn;
+import io.vitess.client.VTGateTx;
+import io.vitess.util.CommonUtils;
+import io.vitess.util.Constants;
+import io.vitess.util.MysqlDefs;
+
 /**
  * Created by harshit.gangal on 23/01/16.
  */
@@ -172,7 +172,7 @@ public void commit() throws SQLException {
         checkAutoCommit(Constants.SQLExceptionMessages.COMMIT_WHEN_AUTO_COMMIT_TRUE);
         try {
             if (isInTransaction()) {
-                Context context = createContext(Constants.CONNECTION_TIMEOUT);
+                Context context = createContext(getTimeout());
                 this.vtGateTx.commit(context, getTwopcEnabled()).checkedGet();
             }
         } finally {
@@ -191,7 +191,7 @@ public void rollback() throws SQLException {
         checkAutoCommit(Constants.SQLExceptionMessages.ROLLBACK_WHEN_AUTO_COMMIT_TRUE);
         try {
             if (isInTransaction()) {
-                Context context = createContext(Constants.CONNECTION_TIMEOUT);
+                Context context = createContext(getTimeout());
                 this.vtGateTx.rollback(context).checkedGet();
             }
         } finally {

diff --git a/java/jdbc/src/main/java/io/vitess/jdbc/VitessStatement.java b/java/jdbc/src/main/java/io/vitess/jdbc/VitessStatement.java
@@ -57,7 +57,7 @@ public class VitessStatement implements Statement {
     protected VitessConnection vitessConnection;
     protected boolean closed;
     protected long resultCount;
-    protected long queryTimeoutInMillis = Constants.DEFAULT_TIMEOUT;
+    protected long queryTimeoutInMillis;
     protected int maxFieldSize = Constants.MAX_BUFFER_SIZE;
     protected int maxRows = 0;
     protected int fetchSize = 0;
@@ -79,6 +79,7 @@ public VitessStatement(VitessConnection vitessConnection) {
     public VitessStatement(VitessConnection vitessConnection, int resultSetType,
         int resultSetConcurrency) {
         this.vitessConnection = vitessConnection;
+        this.queryTimeoutInMillis = vitessConnection.getTimeout();
         this.vitessResultSet = null;
         this.resultSetType = resultSetType;
         this.resultSetConcurrency = resultSetConcurrency;
@@ -269,7 +270,7 @@ public void setQueryTimeout(int seconds) throws SQLException {
                 Constants.SQLExceptionMessages.ILLEGAL_VALUE_FOR + "query timeout");
         }
         this.queryTimeoutInMillis =
-            (0 == seconds) ? Constants.DEFAULT_TIMEOUT : (long) seconds * 1000;
+            (0 == seconds) ? vitessConnection.getTimeout() : (long) seconds * 1000;
     }
 
     /**

diff --git a/java/jdbc/src/main/java/io/vitess/jdbc/VitessVTGateManager.java b/java/jdbc/src/main/java/io/vitess/jdbc/VitessVTGateManager.java
@@ -16,6 +16,13 @@
 
 package io.vitess.jdbc;
 
+import java.io.IOException;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+
 import io.vitess.client.Context;
 import io.vitess.client.RpcClient;
 import io.vitess.client.VTGateConn;
@@ -24,12 +31,6 @@
 import io.vitess.client.grpc.tls.TlsOptions;
 import io.vitess.util.CommonUtils;
 import io.vitess.util.Constants;
-import java.io.IOException;
-import java.sql.SQLException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Random;
-import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Created by naveen.nahata on 24/02/16.
@@ -107,7 +108,7 @@ private static void updateVtGateConnHashMap(String identifier, VitessJDBCUrl.Hos
     private static VTGateConn getVtGateConn(VitessJDBCUrl.HostInfo hostInfo, VitessConnection connection) {
         final String username = connection.getUsername();
         final String keyspace = connection.getKeyspace();
-        final Context context = CommonUtils.createContext(username, Constants.CONNECTION_TIMEOUT);
+        final Context context = CommonUtils.createContext(username,connection.getTimeout());
         RetryingInterceptorConfig retryingConfig = getRetryingInterceptorConfig(connection);
         RpcClient client;
         if (connection.getUseSSL()) {

diff --git a/java/jdbc/src/main/java/io/vitess/util/CommonUtils.java b/java/jdbc/src/main/java/io/vitess/util/CommonUtils.java
@@ -16,9 +16,10 @@
 
 package io.vitess.util;
 
+import org.joda.time.Duration;
+
 import io.vitess.client.Context;
 import io.vitess.proto.Vtrpc;
-import org.joda.time.Duration;
 
 /**
  * Created by naveen.nahata on 24/02/16.
@@ -29,23 +30,24 @@ public class CommonUtils {
      * Create context used to create grpc client and executing query.
      *
      * @param username
-     * @param connectionTimeout
+     * @param timeout
      * @return
      */
-    public static Context createContext(String username, long connectionTimeout) {
-        Context context;
+    public static Context createContext(String username, long timeout) {
+        Context context = Context.getDefault();
         Vtrpc.CallerID callerID = null;
         if (null != username) {
             callerID = Vtrpc.CallerID.newBuilder().setPrincipal(username).build();
         }
+
         if (null != callerID) {
-            context = Context.getDefault().withDeadlineAfter(Duration.millis(connectionTimeout))
-                .withCallerId(callerID);
-        } else {
-            context = Context.getDefault().withDeadlineAfter(Duration.millis(connectionTimeout));
+            context = context.withCallerId(callerID);
         }
+        if (timeout > 0) {
+            context = context.withDeadlineAfter(Duration.millis(timeout));
+        }
+
         return context;
     }
-
 }
 
diff --git a/java/jdbc/src/test/java/io/vitess/jdbc/ConnectionPropertiesTest.java b/java/jdbc/src/test/java/io/vitess/jdbc/ConnectionPropertiesTest.java
@@ -16,21 +16,22 @@
 
 package io.vitess.jdbc;
 
-import io.vitess.proto.Query;
-import io.vitess.proto.Topodata;
-import io.vitess.util.Constants;
-import org.junit.Assert;
-import org.junit.Test;
-import org.mockito.Mockito;
-
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
 import java.util.Arrays;
 import java.util.Properties;
 
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import io.vitess.proto.Query;
+import io.vitess.proto.Topodata;
+import io.vitess.util.Constants;
+
 public class ConnectionPropertiesTest {
 
-    private static final int NUM_PROPS = 31;
+    private static final int NUM_PROPS = 32;
 
     @Test
     public void testReflection() throws Exception {

diff --git a/java/jdbc/src/test/java/io/vitess/jdbc/VitessStatementTest.java b/java/jdbc/src/test/java/io/vitess/jdbc/VitessStatementTest.java
@@ -530,13 +530,23 @@ private void testExecute(int fetchSize, boolean simpleExecute, boolean shouldRun
 
     @Test public void testGetQueryTimeout() throws SQLException {
         VitessConnection mockConn = PowerMockito.mock(VitessConnection.class);
+        Mockito.when(mockConn.getTimeout()).thenReturn((long)Constants.DEFAULT_TIMEOUT);
 
         VitessStatement statement = new VitessStatement(mockConn);
         Assert.assertEquals(30, statement.getQueryTimeout());
     }
 
+    @Test public void testGetQueryTimeoutZeroDefault() throws SQLException {
+        VitessConnection mockConn = PowerMockito.mock(VitessConnection.class);
+        Mockito.when(mockConn.getTimeout()).thenReturn(0L);
+
+        VitessStatement statement = new VitessStatement(mockConn);
+        Assert.assertEquals(0, statement.getQueryTimeout());
+    }
+
     @Test public void testSetQueryTimeout() throws SQLException {
         VitessConnection mockConn = PowerMockito.mock(VitessConnection.class);
+        Mockito.when(mockConn.getTimeout()).thenReturn((long)Constants.DEFAULT_TIMEOUT);
 
         VitessStatement statement = new VitessStatement(mockConn);
 
@@ -550,6 +560,9 @@ private void testExecute(int fetchSize, boolean simpleExecute, boolean shouldRun
         } catch (SQLException ex) {
             Assert.assertEquals("Illegal value for query timeout", ex.getMessage());
         }
+
+        statement.setQueryTimeout(0);
+        Assert.assertEquals(30, statement.getQueryTimeout());
     }
 
     @Test public void testGetWarnings() throws SQLException {