SNOW-1020043: Configure connection and socket timeout (#1628)

src/main/java/net/snowflake/client/core/HttpUtil.java

@@ -20,6 +20,7 @@
 import java.net.Socket;
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.ConcurrentHashMap;
@@ -59,23 +60,27 @@
 import org.apache.http.util.EntityUtils;
 
 public class HttpUtil {
-  static final SFLogger logger = SFLoggerFactory.getLogger(HttpUtil.class);
+  private static final SFLogger logger = SFLoggerFactory.getLogger(HttpUtil.class);
 
   static final int DEFAULT_MAX_CONNECTIONS = 300;
   static final int DEFAULT_MAX_CONNECTIONS_PER_ROUTE = 300;
-  static final int DEFAULT_CONNECTION_TIMEOUT = 60000;
-  static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300000; // ms
+  private static final int DEFAULT_HTTP_CLIENT_CONNECTION_TIMEOUT_IN_MS = 60000;
+  static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT_IN_MS = 300000; // ms
   static final int DEFAULT_TTL = 60; // secs
   static final int DEFAULT_IDLE_CONNECTION_TIMEOUT = 5; // secs
   static final int DEFAULT_DOWNLOADED_CONDITION_TIMEOUT = 3600; // secs
 
   public static final String JDBC_TTL = "net.snowflake.jdbc.ttl";
+  static final String JDBC_CONNECTION_TIMEOUT_IN_MS_PROPERTY =
+      "net.snowflake.jdbc.http_client_connection_timeout_in_ms";
+  static final String JDBC_SOCKET_TIMEOUT_IN_MS_PROPERTY =
+      "net.snowflake.jdbc.http_client_socket_timeout_in_ms";
   public static final String JDBC_MAX_CONNECTIONS_PROPERTY = "net.snowflake.jdbc.max_connections";
   public static final String JDBC_MAX_CONNECTIONS_PER_ROUTE_PROPERTY =
       "net.snowflake.jdbc.max_connections_per_route";
 
   /**
-   * The unique httpClient shared by all connections. This will benefit long- lived clients. Key =
+   * The unique httpClient shared by all connections. This will benefit long-lived clients. Key =
    * proxy host + proxy port + nonProxyHosts, Value = Map of [OCSPMode, HttpClient]
    */
   public static Map<HttpClientSettingsKey, CloseableHttpClient> httpClient =
@@ -101,6 +106,20 @@ public class HttpUtil {
 
   private static boolean socksProxyDisabled = false;
 
+  @SnowflakeJdbcInternalApi
+  public static Duration getConnectionTimeout() {
+    return Duration.ofMillis(
+        convertSystemPropertyToIntValue(
+            JDBC_CONNECTION_TIMEOUT_IN_MS_PROPERTY, DEFAULT_HTTP_CLIENT_CONNECTION_TIMEOUT_IN_MS));
+  }
+
+  @SnowflakeJdbcInternalApi
+  public static Duration getSocketTimeout() {
+    return Duration.ofMillis(
+        convertSystemPropertyToIntValue(
+            JDBC_SOCKET_TIMEOUT_IN_MS_PROPERTY, DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT_IN_MS));
+  }
+
   public static long getDownloadedConditionTimeoutInSeconds() {
     return DEFAULT_DOWNLOADED_CONDITION_TIMEOUT;
   }
@@ -286,9 +305,14 @@ public static CloseableHttpClient buildHttpClient(
       @Nullable HttpClientSettingsKey key, File ocspCacheFile, boolean downloadUnCompressed) {
     // set timeout so that we don't wait forever.
     // Setup the default configuration for all requests on this client
-
     int timeToLive = convertSystemPropertyToIntValue(JDBC_TTL, DEFAULT_TTL);
     logger.debug("time to live in connection pooling manager: {}", timeToLive);
+    long connectTimeout = getConnectionTimeout().toMillis();
+    long socketTimeout = getSocketTimeout().toMillis();
+    logger.debug(
+        "Connect timeout is {} ms and socket timeout is {} for connection pooling manager",
+        connectTimeout,
+        socketTimeout);
 
     // Set proxy settings for DefaultRequestConfig. If current proxy settings are the same as for
     // the last request, keep the current DefaultRequestConfig. If not, build a new
@@ -306,9 +330,9 @@ public static CloseableHttpClient buildHttpClient(
     if (noDefaultRequestConfig || !DefaultRequestConfig.getProxy().equals(proxy)) {
       RequestConfig.Builder builder =
           RequestConfig.custom()
-              .setConnectTimeout(DEFAULT_CONNECTION_TIMEOUT)
-              .setConnectionRequestTimeout(DEFAULT_CONNECTION_TIMEOUT)
-              .setSocketTimeout(DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT);
+              .setConnectTimeout((int) connectTimeout)
+              .setConnectionRequestTimeout((int) connectTimeout)
+              .setSocketTimeout((int) socketTimeout);
       // only set the proxy settings if they are not null
       // but no value has been specified for nonProxyHosts
       // the route planner will determine whether to use a proxy based on nonProxyHosts value.

src/main/java/net/snowflake/client/core/SFLoginInput.java

@@ -7,14 +7,12 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.security.PrivateKey;
+import java.time.Duration;
 import java.util.Map;
 import net.snowflake.client.jdbc.ErrorCode;
 
 /** A class for holding all information required for login */
 public class SFLoginInput {
-  private static int DEFAULT_HTTP_CLIENT_CONNECTION_TIMEOUT = 60000; // millisec
-  private static int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300000; // millisec
-
   private String serverUrl;
   private String databaseName;
   private String schemaName;
@@ -32,8 +30,8 @@ public class SFLoginInput {
   private boolean passcodeInPassword;
   private String passcode;
   private String token;
-  private int connectionTimeout = DEFAULT_HTTP_CLIENT_CONNECTION_TIMEOUT;
-  private int socketTimeout = DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT;
+  private Duration connectionTimeout = HttpUtil.getConnectionTimeout();
+  private Duration socketTimeout = HttpUtil.getSocketTimeout();
   private String appId;
   private String appVersion;
   private String sessionToken;
@@ -216,20 +214,20 @@ public SFLoginInput setToken(String token) {
     return this;
   }
 
-  int getConnectionTimeout() {
+  Duration getConnectionTimeout() {
     return connectionTimeout;
   }
 
-  SFLoginInput setConnectionTimeout(int connectionTimeout) {
+  SFLoginInput setConnectionTimeout(Duration connectionTimeout) {
     this.connectionTimeout = connectionTimeout;
     return this;
   }
 
-  int getSocketTimeout() {
-    return socketTimeout;
+  int getSocketTimeoutInMillis() {
+    return (int) socketTimeout.toMillis();
   }
 
-  SFLoginInput setSocketTimeout(int socketTimeout) {
+  SFLoginInput setSocketTimeout(Duration socketTimeout) {
     this.socketTimeout = socketTimeout;
     return this;
   }

src/main/java/net/snowflake/client/core/SFLoginOutput.java

@@ -4,6 +4,7 @@
 
 package net.snowflake.client.core;
 
+import java.time.Duration;
 import java.util.Map;
 
 /** Login output information including session tokens, database versions */
@@ -16,7 +17,7 @@ public class SFLoginOutput {
   private String databaseVersion;
   private int databaseMajorVersion;
   private int databaseMinorVersion;
-  private int httpClientSocketTimeout;
+  private Duration httpClientSocketTimeout;
   private String sessionDatabase;
   private String sessionSchema;
   private String sessionRole;
@@ -50,7 +51,7 @@ public class SFLoginOutput {
     this.databaseVersion = databaseVersion;
     this.databaseMajorVersion = databaseMajorVersion;
     this.databaseMinorVersion = databaseMinorVersion;
-    this.httpClientSocketTimeout = httpClientSocketTimeout;
+    this.httpClientSocketTimeout = Duration.ofMillis(httpClientSocketTimeout);
     this.sessionDatabase = sessionDatabase;
     this.sessionSchema = sessionSchema;
     this.sessionRole = sessionRole;
@@ -106,7 +107,7 @@ int getDatabaseMinorVersion() {
     return databaseMinorVersion;
   }
 
-  int getHttpClientSocketTimeout() {
+  Duration getHttpClientSocketTimeout() {
     return httpClientSocketTimeout;
   }
 

src/main/java/net/snowflake/client/core/SFSession.java

@@ -14,6 +14,7 @@
 import java.security.PrivateKey;
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
+import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -48,7 +49,9 @@ public class SFSession extends SFBaseSession {
       "CLIENT_STORE_TEMPORARY_CREDENTIAL";
   private static final ObjectMapper mapper = ObjectMapperFactory.getObjectMapper();
   private static final int MAX_SESSION_PARAMETERS = 1000;
-  public static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300000; // millisec
+  // this constant was public - let's not change it
+  public static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT =
+      HttpUtil.DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT_IN_MS;
   private final AtomicInteger sequenceId = new AtomicInteger(0);
   private final List<DriverPropertyInfo> missingProperties = new ArrayList<>();
   // list of active asynchronous queries. Used to see if session should be closed when connection
@@ -85,8 +88,8 @@ public class SFSession extends SFBaseSession {
 
   private int authTimeout = 0;
   private boolean enableCombineDescribe = false;
-  private int httpClientConnectionTimeout = 60000; // milliseconds
-  private int httpClientSocketTimeout = DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT; // milliseconds
+  private final Duration httpClientConnectionTimeout = HttpUtil.getConnectionTimeout();
+  private Duration httpClientSocketTimeout = HttpUtil.getSocketTimeout();
   // whether we try to simulate a socket timeout (a default value of 0 means
   // no simulation). The value is in milliseconds
   private int injectSocketTimeout = 0;
@@ -188,7 +191,12 @@ private JsonNode getQueryMetadata(String queryID) throws SQLException {
         get.setHeader("Authorization", "Snowflake Token=\"" + this.sessionToken + "\"");
         response =
             HttpUtil.executeGeneralRequest(
-                get, loginTimeout, authTimeout, httpClientSocketTimeout, 0, getHttpClientKey());
+                get,
+                loginTimeout,
+                authTimeout,
+                (int) httpClientSocketTimeout.toMillis(),
+                0,
+                getHttpClientKey());
         jsonNode = OBJECT_MAPPER.readTree(response);
       } catch (Exception e) {
         throw new SnowflakeSQLLoggedException(
@@ -915,7 +923,7 @@ protected void heartbeat() throws SFException, SQLException {
                 postRequest,
                 SF_HEARTBEAT_TIMEOUT,
                 authTimeout,
-                httpClientSocketTimeout,
+                (int) httpClientSocketTimeout.toMillis(),
                 0,
                 getHttpClientKey());
 
@@ -985,11 +993,11 @@ public int getAuthTimeout() {
   }
 
   public int getHttpClientSocketTimeout() {
-    return httpClientSocketTimeout;
+    return (int) httpClientSocketTimeout.toMillis();
   }
 
   public int getHttpClientConnectionTimeout() {
-    return httpClientConnectionTimeout;
+    return (int) httpClientConnectionTimeout.toMillis();
   }
 
   public boolean isClosed() {

src/main/java/net/snowflake/client/core/SessionUtil.java

@@ -354,7 +354,7 @@ private static SFLoginOutput newSession(
     int databaseMinorVersion = 0;
     String newClientForUpgrade;
     int healthCheckInterval = DEFAULT_HEALTH_CHECK_INTERVAL;
-    int httpClientSocketTimeout = loginInput.getSocketTimeout();
+    int httpClientSocketTimeout = loginInput.getSocketTimeoutInMillis();
     final ClientAuthnDTO.AuthenticatorType authenticatorType = getAuthenticator(loginInput);
     Map<String, Object> commonParams;
 
@@ -623,7 +623,7 @@ private static SFLoginOutput newSession(
       String theString = null;
 
       int leftRetryTimeout = loginInput.getLoginTimeout();
-      int leftsocketTimeout = loginInput.getSocketTimeout();
+      int leftsocketTimeout = loginInput.getSocketTimeoutInMillis();
       int retryCount = 0;
 
       while (true) {
@@ -664,7 +664,7 @@ private static SFLoginOutput newSession(
               // auth timeout within socket timeout is thrown without backoff,
               // and we need to update time remained in socket timeout here to control the
               // the actual socket timeout from customer setting.
-              if (loginInput.getSocketTimeout() > 0) {
+              if (loginInput.getSocketTimeoutInMillis() > 0) {
                 if (ex.issocketTimeoutNoBackoff()) {
                   if (leftsocketTimeout > elapsedSeconds) {
                     leftsocketTimeout -= elapsedSeconds;
@@ -673,7 +673,7 @@ private static SFLoginOutput newSession(
                   }
                 } else {
                   // reset curl timeout for retry with backoff.
-                  leftsocketTimeout = loginInput.getSocketTimeout();
+                  leftsocketTimeout = loginInput.getSocketTimeoutInMillis();
                 }
               }
 
@@ -782,11 +782,11 @@ private static SFLoginOutput newSession(
       if (healthCheckIntervalFromGS > 0 && healthCheckIntervalFromGS != healthCheckInterval) {
         // add health check interval to socket timeout
         httpClientSocketTimeout =
-            loginInput.getSocketTimeout() + (healthCheckIntervalFromGS * 1000);
+            loginInput.getSocketTimeoutInMillis() + (healthCheckIntervalFromGS * 1000);
 
         final RequestConfig requestConfig =
             RequestConfig.copy(HttpUtil.getRequestConfigWithoutCookies())
-                .setConnectTimeout(loginInput.getConnectionTimeout())
+                .setConnectTimeout((int) loginInput.getConnectionTimeout().toMillis())
                 .setSocketTimeout(httpClientSocketTimeout)
                 .build();
 
@@ -961,7 +961,7 @@ private static SFLoginOutput tokenRequest(SFLoginInput loginInput, TokenRequestT
               postRequest,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
 
@@ -1054,7 +1054,7 @@ static void closeSession(SFLoginInput loginInput) throws SFException, SnowflakeS
               postRequest,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
 
@@ -1120,7 +1120,7 @@ private static String federatedFlowStep4(
               httpGet,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
 
@@ -1194,7 +1194,7 @@ private static String federatedFlowStep3(SFLoginInput loginInput, String tokenUr
               postRequest,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               0,
               null,
@@ -1285,7 +1285,7 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf
               postRequest,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
       logger.debug("authenticator-request response: {}", gsResponse);

src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java

@@ -187,7 +187,7 @@ private String getSSOUrl(int port) throws SFException, SnowflakeSQLException {
               postRequest,
               loginInput.getLoginTimeout(),
               loginInput.getAuthTimeout(),
-              loginInput.getSocketTimeout(),
+              loginInput.getSocketTimeoutInMillis(),
               0,
               loginInput.getHttpClientSettingsKey());
 

src/main/java/net/snowflake/client/jdbc/RestRequest.java

@@ -242,7 +242,8 @@ public static CloseableHttpResponse execute(
 
         savedEx = ex;
         // if the request took more than 5 min (socket timeout) log an error
-        if ((System.currentTimeMillis() - startTimePerRequest) > 300000) {
+        if ((System.currentTimeMillis() - startTimePerRequest)
+            > HttpUtil.getSocketTimeout().toMillis()) {
           logger.warn(
               "HTTP request took longer than 5 min: {} sec",
               (System.currentTimeMillis() - startTimePerRequest) / 1000);

src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java

@@ -584,7 +584,7 @@ public void uploadWithPresignedUrlWithoutConnection(
       uploadWithPresignedUrl(
           networkTimeoutInMilli,
           0, // auth timeout
-          SFSession.DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT, // socket timeout
+          (int) HttpUtil.getSocketTimeout().toMillis(),
           meta.getContentEncoding(),
           meta.getUserMetadata(),
           uploadStreamInfo.left,

src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java

@@ -3,8 +3,6 @@
  */
 package net.snowflake.client.jdbc.telemetry;
 
-import static net.snowflake.client.core.SFSession.DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT;
-
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
@@ -361,7 +359,7 @@ private boolean sendBatch() throws IOException {
                     post,
                     TELEMETRY_HTTP_RETRY_TIMEOUT_IN_SEC,
                     0,
-                    DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT,
+                    (int) HttpUtil.getSocketTimeout().toMillis(),
                     0,
                     this.httpClient)
                 : HttpUtil.executeGeneralRequest(