From ff55bfa3d342380d88f6a79ff28ca224d5a1a55e Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Tue, 31 Oct 2023 04:36:24 -0700 Subject: [PATCH 01/10] Added new rules for retry for login endpoints --- .../net/snowflake/client/core/SFSession.java | 4 +-- .../snowflake/client/core/SessionUtil.java | 32 +++++++++++++++++++ .../snowflake/client/jdbc/RestRequest.java | 26 +++++++++++++-- .../util/DecorrelatedJitterBackoff.java | 12 +++++++ 4 files changed, 70 insertions(+), 4 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java index e0ddcc7a8..4316e8420 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -70,9 +70,9 @@ public class SFSession extends SFBaseSession { * Amount of seconds a user is willing to tolerate for establishing the connection with database. * In our case, it means the first login request to get authorization token. * - *

Default:60 seconds + *

Default:300 seconds */ - private int loginTimeout = 60; + private int loginTimeout = 300; /** * Amount of milliseconds a user is willing to tolerate for network related issues (e.g. HTTP * 503/504) or database transient issues (e.g. GS not responding) diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java index f50586b42..4390cef99 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtil.java +++ b/src/main/java/net/snowflake/client/core/SessionUtil.java @@ -33,6 +33,7 @@ import org.apache.http.client.config.RequestConfig; import org.apache.http.client.methods.HttpGet; import org.apache.http.client.methods.HttpPost; +import org.apache.http.client.methods.HttpRequestBase; import org.apache.http.client.utils.URIBuilder; import org.apache.http.entity.StringEntity; import org.apache.http.message.BasicHeader; @@ -112,6 +113,14 @@ public class SessionUtil { static final String SF_HEADER_SERVICE_NAME = "X-Snowflake-Service"; + public static final String SF_HEADER_CLIENT_APP_ID = "CLIENT_APP_ID"; + + public static final String SF_HEADER_CLIENT_APP_VERSION = "CLIENT_APP_VERSION"; + + private static final String SF_DRIVER_NAME = "Snowflake"; + + private static final String SF_DRIVER_VERSION = SnowflakeDriver.implementVersion; + private static final String ID_TOKEN_AUTHENTICATOR = "ID_TOKEN"; private static final String NO_QUERY_ID = ""; @@ -592,6 +601,10 @@ private static SFLoginOutput newSession( HttpUtil.applyAdditionalHeadersForSnowsight( postRequest, loginInput.getAdditionalHttpHeadersForSnowsight()); + // Add headers for driver name and version + postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, SF_DRIVER_NAME); + postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, SF_DRIVER_VERSION); + // attach the login info json body to the post request StringEntity input = new StringEntity(json, StandardCharsets.UTF_8); input.setContentType("application/json"); @@ -1614,4 +1627,23 @@ public static String generateJWTToken( privateKey, privateKeyFile, privateKeyFilePwd, accountName, userName); return s.issueJwtToken(); } + + /** + * Helper method to check if the request path is a login/auth request to use for retry strategy. + * + * @param request the post request + * @return true if this is a login/auth request, false otherwise + */ + public static boolean isLoginRequest(HttpRequestBase request) { + URI requestURI = request.getURI(); + String requestPath = requestURI.getPath(); + if (requestPath != null) { + if (requestPath.equals(SF_PATH_LOGIN_REQUEST) + || requestPath.equals(SF_PATH_AUTHENTICATOR_REQUEST) + || requestPath.equals(SF_PATH_TOKEN_REQUEST)) { + return true; + } + } + return false; + } } diff --git a/src/main/java/net/snowflake/client/jdbc/RestRequest.java b/src/main/java/net/snowflake/client/jdbc/RestRequest.java index 52d02768c..30b55e76b 100644 --- a/src/main/java/net/snowflake/client/jdbc/RestRequest.java +++ b/src/main/java/net/snowflake/client/jdbc/RestRequest.java @@ -42,6 +42,9 @@ public class RestRequest { // min backoff in milli before we retry due to transient issues private static final long minBackoffInMilli = 1000; + // min backoff in milli for login/auth requests before we retry + private static final long minLoginBackoffInMilli = 4000; + // max backoff in milli before we retry due to transient issues // we double the backoff after each retry till we reach the max backoff private static final long maxBackoffInMilli = 16000; @@ -132,6 +135,9 @@ public static CloseableHttpResponse execute( // when there are transient network/GS issues. long startTimePerRequest = startTime; + // Used to indicate that this is a login/auth request and will be using the new retry strategy. + boolean isLoginRequest = SessionUtil.isLoginRequest(httpRequest); + // total elapsed time due to transient issues. long elapsedMilliForTransientIssues = 0; @@ -139,7 +145,12 @@ public static CloseableHttpResponse execute( long retryTimeoutInMilliseconds = retryTimeout * 1000; // amount of time to wait for backing off before retry - long backoffInMilli = minBackoffInMilli; + long backoffInMilli; + if (isLoginRequest) { + backoffInMilli = minLoginBackoffInMilli; + } else { + backoffInMilli = minBackoffInMilli; + } // auth timeout (ms) long authTimeoutInMilli = authTimeout * 1000; @@ -417,7 +428,18 @@ public static CloseableHttpResponse execute( logger.debug("sleeping in {}(ms)", backoffInMilli); Thread.sleep(backoffInMilli); elapsedMilliForTransientIssues += backoffInMilli; - backoffInMilli = backoff.nextSleepTime(backoffInMilli); + if (isLoginRequest) { + backoffInMilli = backoff.getJitterForLogin(backoffInMilli); + } else { + backoffInMilli = backoff.nextSleepTime(backoffInMilli); + } + if (retryTimeoutInMilliseconds > 0 + && (elapsedMilliForTransientIssues + backoffInMilli) > retryTimeoutInMilliseconds) { + // If the timeout will be reached before the next backoff, just use the remaining time. + backoffInMilli = + Math.min( + backoffInMilli, retryTimeoutInMilliseconds - elapsedMilliForTransientIssues); + } } catch (InterruptedException ex1) { logger.debug("Backoff sleep before retrying login got interrupted", false); } diff --git a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java index a7f98fe01..3b34447ba 100644 --- a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java +++ b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java @@ -1,5 +1,6 @@ package net.snowflake.client.util; +import java.util.Random; import java.util.concurrent.ThreadLocalRandom; /** @@ -19,4 +20,15 @@ public DecorrelatedJitterBackoff(long base, long cap) { public long nextSleepTime(long sleep) { return Math.min(cap, ThreadLocalRandom.current().nextLong(base, sleep * 3)); } + + public long getJitterForLogin(long currentTime) { + int mulitplicationFactor = chooseRandom(-1, 1); + long jitter = (long) (mulitplicationFactor * currentTime * 0.5); + return jitter; + } + + private int chooseRandom(int min, int max) { + Random random = new Random(); + return random.nextInt(max - min) + min; + } } From ff4939318387be297bbc9391bdf690c9ebe7d797 Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Tue, 31 Oct 2023 04:42:51 -0700 Subject: [PATCH 02/10] fix check-style --- src/main/java/net/snowflake/client/jdbc/RestRequest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/java/net/snowflake/client/jdbc/RestRequest.java b/src/main/java/net/snowflake/client/jdbc/RestRequest.java index 30b55e76b..ba06ec411 100644 --- a/src/main/java/net/snowflake/client/jdbc/RestRequest.java +++ b/src/main/java/net/snowflake/client/jdbc/RestRequest.java @@ -435,7 +435,8 @@ public static CloseableHttpResponse execute( } if (retryTimeoutInMilliseconds > 0 && (elapsedMilliForTransientIssues + backoffInMilli) > retryTimeoutInMilliseconds) { - // If the timeout will be reached before the next backoff, just use the remaining time. + // If the timeout will be reached before the next backoff, just use the remaining + // time. backoffInMilli = Math.min( backoffInMilli, retryTimeoutInMilliseconds - elapsedMilliForTransientIssues); From d73b29c45235a7a3e563ed6c6a4f7c8edcb2664a Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Tue, 31 Oct 2023 07:29:34 -0700 Subject: [PATCH 03/10] Added tests --- .../client/core/SessionUtilTest.java | 46 ++++++++++- .../client/jdbc/ConnectionLatestIT.java | 4 +- .../client/jdbc/RestRequestTest.java | 80 ++++++++++++++++++- 3 files changed, 126 insertions(+), 4 deletions(-) diff --git a/src/test/java/net/snowflake/client/core/SessionUtilTest.java b/src/test/java/net/snowflake/client/core/SessionUtilTest.java index aadcaa791..1846cc1d1 100644 --- a/src/test/java/net/snowflake/client/core/SessionUtilTest.java +++ b/src/test/java/net/snowflake/client/core/SessionUtilTest.java @@ -5,12 +5,18 @@ package net.snowflake.client.core; import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.Assert.assertEquals; +import static org.junit.Assert.*; import com.fasterxml.jackson.databind.node.BooleanNode; +import java.net.URI; +import java.net.URISyntaxException; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import net.snowflake.client.jdbc.MockConnectionTest; +import org.apache.http.client.methods.HttpPost; +import org.apache.http.client.utils.URIBuilder; import org.junit.Test; public class SessionUtilTest { @@ -83,4 +89,42 @@ public void testConvertSystemPropertyToIntValue() { HttpUtil.JDBC_MAX_CONNECTIONS_PER_ROUTE_PROPERTY, HttpUtil.DEFAULT_MAX_CONNECTIONS_PER_ROUTE)); } + + @Test + public void testIsLoginRequest() { + List testCases = new ArrayList(); + testCases.add("/session/v1/login-request"); + testCases.add("/session/token-request"); + testCases.add("/session/authenticator-request"); + + for (String testCase : testCases) { + try { + URIBuilder uriBuilder = new URIBuilder("https://test.snowflakecomputing.com"); + uriBuilder.setPath(testCase); + URI uri = uriBuilder.build(); + HttpPost postRequest = new HttpPost(uri); + assertTrue(SessionUtil.isLoginRequest(postRequest)); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + } + + @Test + public void testIsLoginRequestInvalidURIPath() { + List testCases = new ArrayList(); + testCases.add("/session/not-a-real-path"); + + for (String testCase : testCases) { + try { + URIBuilder uriBuilder = new URIBuilder("https://test.snowflakecomputing.com"); + uriBuilder.setPath(testCase); + URI uri = uriBuilder.build(); + HttpPost postRequest = new HttpPost(uri); + assertFalse(SessionUtil.isLoginRequest(postRequest)); + } catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + } } diff --git a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java index 58c03c8f4..44f120026 100644 --- a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java +++ b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java @@ -554,7 +554,7 @@ public void testWrongHostNameTimeout() throws InterruptedException { equalTo(ErrorCode.NETWORK_ERROR.getMessageCode())); conEnd = System.currentTimeMillis(); - assertThat("Login time out not taking effective", conEnd - connStart < 60000); + assertThat("Login time out not taking effective", conEnd - connStart < 300000); Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS); if (TelemetryService.getInstance().isDeploymentEnabled()) { @@ -595,7 +595,7 @@ public void testHttpsLoginTimeoutWithSSL() throws InterruptedException { equalTo(ErrorCode.NETWORK_ERROR.getMessageCode())); conEnd = System.currentTimeMillis(); - assertThat("Login time out not taking effective", conEnd - connStart < 60000); + assertThat("Login time out not taking effective", conEnd - connStart < 300000); Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS); if (TelemetryService.getInstance().isDeploymentEnabled()) { assertThat( diff --git a/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java b/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java index 5509ca2c7..dc00d053a 100644 --- a/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java +++ b/src/test/java/net/snowflake/client/jdbc/RestRequestTest.java @@ -29,7 +29,7 @@ /** RestRequest unit tests. */ public class RestRequestTest { - static final int DEFAULT_CONNECTION_TIMEOUT = 60000; + static final int DEFAULT_CONNECTION_TIMEOUT = 300000; static final int DEFAULT_HTTP_CLIENT_SOCKET_TIMEOUT = 300000; // ms private CloseableHttpResponse retryResponse() { @@ -42,6 +42,16 @@ private CloseableHttpResponse retryResponse() { return retryResponse; } + private CloseableHttpResponse retryLoginResponse() { + StatusLine retryStatusLine = mock(StatusLine.class); + when(retryStatusLine.getStatusCode()).thenReturn(429); + + CloseableHttpResponse retryResponse = mock(CloseableHttpResponse.class); + when(retryResponse.getStatusLine()).thenReturn(retryStatusLine); + + return retryResponse; + } + private CloseableHttpResponse successResponse() { StatusLine successStatusLine = mock(StatusLine.class); when(successStatusLine.getStatusCode()).thenReturn(200); @@ -457,6 +467,74 @@ public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwabl } } + @Test(expected = SnowflakeSQLException.class) + public void testLoginMaxRetries() throws IOException, SnowflakeSQLException { + boolean telemetryEnabled = TelemetryService.getInstance().isEnabled(); + + CloseableHttpClient client = mock(CloseableHttpClient.class); + when(client.execute(any(HttpUriRequest.class))) + .thenAnswer( + new Answer() { + int callCount = 0; + + @Override + public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable { + callCount += 1; + if (callCount >= 4) { + return retryLoginResponse(); + } else { + return socketTimeoutResponse(); + } + } + }); + + try { + TelemetryService.disable(); + execute(client, "/session/v1/login-request", 0, 0, 0, true, false, 1); + fail("testMaxRetries"); + } finally { + if (telemetryEnabled) { + TelemetryService.enable(); + } else { + TelemetryService.disable(); + } + } + } + + @Test(expected = SnowflakeSQLException.class) + public void testLoginTimeout() throws IOException, SnowflakeSQLException { + boolean telemetryEnabled = TelemetryService.getInstance().isEnabled(); + + CloseableHttpClient client = mock(CloseableHttpClient.class); + when(client.execute(any(HttpUriRequest.class))) + .thenAnswer( + new Answer() { + int callCount = 0; + + @Override + public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwable { + callCount += 1; + if (callCount >= 4) { + return retryLoginResponse(); + } else { + return socketTimeoutResponse(); + } + } + }); + + try { + TelemetryService.disable(); + execute(client, "/session/v1/login-request", 1, 0, 0, true, false, 10); + fail("testMaxRetries"); + } finally { + if (telemetryEnabled) { + TelemetryService.enable(); + } else { + TelemetryService.disable(); + } + } + } + @Test public void testMaxRetriesWithSuccessfulResponse() throws IOException { boolean telemetryEnabled = TelemetryService.getInstance().isEnabled(); From 384aa562a3c080f4d239578a5d9cd1a14d7a12de Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Tue, 31 Oct 2023 13:58:31 -0700 Subject: [PATCH 04/10] modified jitter --- .../java/net/snowflake/client/core/SessionUtil.java | 12 ++++++++++-- .../client/util/DecorrelatedJitterBackoff.java | 9 ++++----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java index 4390cef99..1e071d73f 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtil.java +++ b/src/main/java/net/snowflake/client/core/SessionUtil.java @@ -602,8 +602,8 @@ private static SFLoginOutput newSession( postRequest, loginInput.getAdditionalHttpHeadersForSnowsight()); // Add headers for driver name and version - postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, SF_DRIVER_NAME); - postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, SF_DRIVER_VERSION); + postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId()); + postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion()); // attach the login info json body to the post request StringEntity input = new StringEntity(json, StandardCharsets.UTF_8); @@ -915,6 +915,10 @@ private static SFLoginOutput tokenRequest(SFLoginInput loginInput, TokenRequestT postRequest = new HttpPost(uriBuilder.build()); + // Add headers for driver name and version + postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId()); + postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion()); + // Add custom headers before adding common headers HttpUtil.applyAdditionalHeadersForSnowsight( postRequest, loginInput.getAdditionalHttpHeadersForSnowsight()); @@ -1272,6 +1276,10 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf postRequest.setEntity(input); postRequest.addHeader("accept", "application/json"); + // Add headers for driver name and version + postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId()); + postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion()); + final String gsResponse = HttpUtil.executeGeneralRequest( postRequest, diff --git a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java index 3b34447ba..4cfdbc7ae 100644 --- a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java +++ b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java @@ -22,13 +22,12 @@ public long nextSleepTime(long sleep) { } public long getJitterForLogin(long currentTime) { - int mulitplicationFactor = chooseRandom(-1, 1); - long jitter = (long) (mulitplicationFactor * currentTime * 0.5); + double multiplicationFactor = chooseRandom(-1, 1); + long jitter = (long) (multiplicationFactor * currentTime * 0.5); return jitter; } - private int chooseRandom(int min, int max) { - Random random = new Random(); - return random.nextInt(max - min) + min; + private double chooseRandom(int min, int max) { + return min + (Math.random() * (max - min)); } } From cffa2fd48d36399d5a4f37e55fa47e4387ecc71a Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Tue, 31 Oct 2023 14:03:39 -0700 Subject: [PATCH 05/10] check-style update --- .../net/snowflake/client/util/DecorrelatedJitterBackoff.java | 1 - 1 file changed, 1 deletion(-) diff --git a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java index 4cfdbc7ae..a988e7206 100644 --- a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java +++ b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java @@ -1,6 +1,5 @@ package net.snowflake.client.util; -import java.util.Random; import java.util.concurrent.ThreadLocalRandom; /** From 52832bc268bcf09203c9e0d51e51bd027d0bbe8c Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Fri, 3 Nov 2023 02:55:22 -0700 Subject: [PATCH 06/10] Updated implementation for jitter and added new connection param retryTimeout --- .../snowflake/client/core/SFLoginInput.java | 10 +++++++++ .../net/snowflake/client/core/SFSession.java | 22 ++++++++++++++++++- .../client/core/SFSessionProperty.java | 4 +++- .../snowflake/client/core/SessionUtil.java | 6 +++++ .../snowflake/client/jdbc/RestRequest.java | 17 ++++++-------- .../util/DecorrelatedJitterBackoff.java | 2 +- 6 files changed, 48 insertions(+), 13 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/src/main/java/net/snowflake/client/core/SFLoginInput.java index 35df50a7c..27e452ae5 100644 --- a/src/main/java/net/snowflake/client/core/SFLoginInput.java +++ b/src/main/java/net/snowflake/client/core/SFLoginInput.java @@ -25,6 +25,7 @@ public class SFLoginInput { private String oktaUserName; private String accountName; private int loginTimeout = -1; // default is invalid + private int retryTimeout = 0; private int authTimeout = 0; private String userName; private String password; @@ -144,6 +145,15 @@ SFLoginInput setLoginTimeout(int loginTimeout) { return this; } + int getRetryTimeout() { + return retryTimeout; + } + + SFLoginInput setRetryTimeout(int retryTimeout) { + this.retryTimeout = retryTimeout; + return this; + } + int getAuthTimeout() { return authTimeout; } diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java index 4316e8420..bd5bee1e3 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -108,6 +108,13 @@ public class SFSession extends SFBaseSession { // Max retries for outgoing http requests. private int maxHttpRetries = 7; + /** + * Retry timeout in seconds. Cannot be less than 300. + * + *

Default: 300 + */ + private int retryTimeout = 300; + // This constructor is used only by tests with no real connection. // For real connections, the other constructor is always used. @VisibleForTesting @@ -369,6 +376,15 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro } break; + case RETRY_TIMEOUT: + if (propertyValue != null) { + int timeoutValue = (Integer) propertyValue; + if (timeoutValue >= 300) { + retryTimeout = timeoutValue; + } + } + break; + default: break; } @@ -405,7 +421,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException { "input: server={}, account={}, user={}, password={}, role={}, database={}, schema={}," + " warehouse={}, validate_default_parameters={}, authenticator={}, ocsp_mode={}," + " passcode_in_password={}, passcode={}, private_key={}, disable_socks_proxy={}," - + " application={}, app_id={}, app_version={}, login_timeout={}, network_timeout={}," + + " application={}, app_id={}, app_version={}, login_timeout={}, retry_timeout={}, network_timeout={}," + " query_timeout={}, tracing={}, private_key_file={}, private_key_file_pwd={}." + " session_parameters: client_store_temporary_credential={}, gzip_disabled={}", connectionPropertiesMap.get(SFSessionProperty.SERVER_URL), @@ -433,6 +449,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException { connectionPropertiesMap.get(SFSessionProperty.APP_ID), connectionPropertiesMap.get(SFSessionProperty.APP_VERSION), connectionPropertiesMap.get(SFSessionProperty.LOGIN_TIMEOUT), + connectionPropertiesMap.get(SFSessionProperty.RETRY_TIMEOUT), connectionPropertiesMap.get(SFSessionProperty.NETWORK_TIMEOUT), connectionPropertiesMap.get(SFSessionProperty.QUERY_TIMEOUT), connectionPropertiesMap.get(SFSessionProperty.TRACING), @@ -471,6 +488,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException { .setOKTAUserName((String) connectionPropertiesMap.get(SFSessionProperty.OKTA_USERNAME)) .setAccountName((String) connectionPropertiesMap.get(SFSessionProperty.ACCOUNT)) .setLoginTimeout(loginTimeout) + .setRetryTimeout(retryTimeout) .setAuthTimeout(authTimeout) .setUserName((String) connectionPropertiesMap.get(SFSessionProperty.USER)) .setPassword((String) connectionPropertiesMap.get(SFSessionProperty.PASSWORD)) @@ -652,6 +670,7 @@ synchronized void renewSession(String prevSessionToken) .setIdToken(idToken) .setMfaToken(mfaToken) .setLoginTimeout(loginTimeout) + .setRetryTimeout(retryTimeout) .setDatabaseName(getDatabase()) .setSchemaName(getSchema()) .setRole(getRole()) @@ -696,6 +715,7 @@ public void close() throws SFException, SnowflakeSQLException { .setServerUrl(getServerUrl()) .setSessionToken(sessionToken) .setLoginTimeout(loginTimeout) + .setRetryTimeout(retryTimeout) .setOCSPMode(getOCSPMode()) .setHttpClientSettingsKey(getHttpClientKey()); diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java index 8609cbdc5..97a921c01 100644 --- a/src/main/java/net/snowflake/client/core/SFSessionProperty.java +++ b/src/main/java/net/snowflake/client/core/SFSessionProperty.java @@ -73,7 +73,9 @@ public enum SFSessionProperty { ENABLE_PUT_GET("enablePutGet", false, Boolean.class), - PUT_GET_MAX_RETRIES("putGetMaxRetries", false, Integer.class); + PUT_GET_MAX_RETRIES("putGetMaxRetries", false, Integer.class), + + RETRY_TIMEOUT("retryTimeout", false, Integer.class); // property key in string private String propertyKey; diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java index 1e071d73f..c1cc6a6e1 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtil.java +++ b/src/main/java/net/snowflake/client/core/SessionUtil.java @@ -622,7 +622,13 @@ private static SFLoginOutput newSession( setServiceNameHeader(loginInput, postRequest); String theString = null; + + // We want to choose the smaller of the two values between retryTimeout and loginTimeout int leftRetryTimeout = loginInput.getLoginTimeout(); + if (leftRetryTimeout > loginInput.getRetryTimeout()) { + leftRetryTimeout = loginInput.getRetryTimeout(); + } + int leftsocketTimeout = loginInput.getSocketTimeout(); int retryCount = 0; diff --git a/src/main/java/net/snowflake/client/jdbc/RestRequest.java b/src/main/java/net/snowflake/client/jdbc/RestRequest.java index ba06ec411..3f6b0884a 100644 --- a/src/main/java/net/snowflake/client/jdbc/RestRequest.java +++ b/src/main/java/net/snowflake/client/jdbc/RestRequest.java @@ -42,9 +42,6 @@ public class RestRequest { // min backoff in milli before we retry due to transient issues private static final long minBackoffInMilli = 1000; - // min backoff in milli for login/auth requests before we retry - private static final long minLoginBackoffInMilli = 4000; - // max backoff in milli before we retry due to transient issues // we double the backoff after each retry till we reach the max backoff private static final long maxBackoffInMilli = 16000; @@ -145,12 +142,7 @@ public static CloseableHttpResponse execute( long retryTimeoutInMilliseconds = retryTimeout * 1000; // amount of time to wait for backing off before retry - long backoffInMilli; - if (isLoginRequest) { - backoffInMilli = minLoginBackoffInMilli; - } else { - backoffInMilli = minBackoffInMilli; - } + long backoffInMilli = minBackoffInMilli; // auth timeout (ms) long authTimeoutInMilli = authTimeout * 1000; @@ -429,7 +421,12 @@ public static CloseableHttpResponse execute( Thread.sleep(backoffInMilli); elapsedMilliForTransientIssues += backoffInMilli; if (isLoginRequest) { - backoffInMilli = backoff.getJitterForLogin(backoffInMilli); + long jitteredBackoffInMilli = backoff.getJitterForLogin(backoffInMilli); + backoffInMilli = + (long) + backoff.chooseRandom( + jitteredBackoffInMilli + backoffInMilli, + Math.pow(retryCount, 2) + jitteredBackoffInMilli); } else { backoffInMilli = backoff.nextSleepTime(backoffInMilli); } diff --git a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java index a988e7206..8c73a817a 100644 --- a/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java +++ b/src/main/java/net/snowflake/client/util/DecorrelatedJitterBackoff.java @@ -26,7 +26,7 @@ public long getJitterForLogin(long currentTime) { return jitter; } - private double chooseRandom(int min, int max) { + public double chooseRandom(double min, double max) { return min + (Math.random() * (max - min)); } } From 759c5e34d12331393b7976a8c03c3d08466f679d Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Fri, 3 Nov 2023 03:09:20 -0700 Subject: [PATCH 07/10] - adding timeout for all login endpoints --- .../snowflake/client/core/SessionUtil.java | 20 +++++++++++++------ .../core/SessionUtilExternalBrowser.java | 8 +++++++- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java index c1cc6a6e1..2bc2a17da 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtil.java +++ b/src/main/java/net/snowflake/client/core/SessionUtil.java @@ -117,10 +117,6 @@ public class SessionUtil { public static final String SF_HEADER_CLIENT_APP_VERSION = "CLIENT_APP_VERSION"; - private static final String SF_DRIVER_NAME = "Snowflake"; - - private static final String SF_DRIVER_VERSION = SnowflakeDriver.implementVersion; - private static final String ID_TOKEN_AUTHENTICATOR = "ID_TOKEN"; private static final String NO_QUERY_ID = ""; @@ -962,10 +958,16 @@ private static SFLoginOutput tokenRequest(SFLoginInput loginInput, TokenRequestT (ArgSupplier) () -> loginInput.getSessionToken() != null ? "******" : null, (ArgSupplier) () -> loginInput.getMasterToken() != null ? "******" : null); + // We want to choose the smaller of the two values between retryTimeout and loginTimeout + int loginRetryTimeout = loginInput.getLoginTimeout(); + if (loginRetryTimeout > loginInput.getRetryTimeout()) { + loginRetryTimeout = loginInput.getRetryTimeout(); + } + String theString = HttpUtil.executeGeneralRequest( postRequest, - loginInput.getLoginTimeout(), + loginRetryTimeout, loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, @@ -1286,10 +1288,16 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId()); postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion()); + // We want to choose the smaller of the two values between retryTimeout and loginTimeout + int loginRetryTimeout = loginInput.getLoginTimeout(); + if (loginRetryTimeout > loginInput.getRetryTimeout()) { + loginRetryTimeout = loginInput.getRetryTimeout(); + } + final String gsResponse = HttpUtil.executeGeneralRequest( postRequest, - loginInput.getLoginTimeout(), + loginRetryTimeout, loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, diff --git a/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java b/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java index b18d1ebcb..b7194ba0a 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java +++ b/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java @@ -181,10 +181,16 @@ private String getSSOUrl(int port) throws SFException, SnowflakeSQLException { postRequest.addHeader("accept", "application/json"); + // We want to choose the smaller of the two values between retryTimeout and loginTimeout + int loginRetryTimeout = loginInput.getLoginTimeout(); + if (loginRetryTimeout > loginInput.getRetryTimeout()) { + loginRetryTimeout = loginInput.getRetryTimeout(); + } + String theString = HttpUtil.executeGeneralRequest( postRequest, - loginInput.getLoginTimeout(), + loginRetryTimeout, loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, From 5dab0bc05d1ecd5a3701b66adb6e1a338229b673 Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Fri, 3 Nov 2023 14:57:36 -0700 Subject: [PATCH 08/10] code review changes --- .../snowflake/client/core/SFLoginInput.java | 10 ++++++-- .../snowflake/client/core/SessionUtil.java | 23 +++---------------- .../core/SessionUtilExternalBrowser.java | 8 +------ .../snowflake/client/jdbc/RestRequest.java | 4 ++-- .../client/core/SessionUtilTest.java | 4 ++-- 5 files changed, 16 insertions(+), 33 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/src/main/java/net/snowflake/client/core/SFLoginInput.java index 27e452ae5..d48e3d185 100644 --- a/src/main/java/net/snowflake/client/core/SFLoginInput.java +++ b/src/main/java/net/snowflake/client/core/SFLoginInput.java @@ -25,7 +25,7 @@ public class SFLoginInput { private String oktaUserName; private String accountName; private int loginTimeout = -1; // default is invalid - private int retryTimeout = 0; + private int retryTimeout = 300; private int authTimeout = 0; private String userName; private String password; @@ -140,8 +140,14 @@ int getLoginTimeout() { return loginTimeout; } + // We want to choose the smaller of the two values between retryTimeout and loginTimeout for the + // new retry strategy. SFLoginInput setLoginTimeout(int loginTimeout) { - this.loginTimeout = loginTimeout; + if (loginTimeout > retryTimeout) { + this.loginTimeout = retryTimeout; + } else { + this.loginTimeout = loginTimeout; + } return this; } diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java index 2bc2a17da..18710aaf2 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtil.java +++ b/src/main/java/net/snowflake/client/core/SessionUtil.java @@ -619,12 +619,7 @@ private static SFLoginOutput newSession( String theString = null; - // We want to choose the smaller of the two values between retryTimeout and loginTimeout int leftRetryTimeout = loginInput.getLoginTimeout(); - if (leftRetryTimeout > loginInput.getRetryTimeout()) { - leftRetryTimeout = loginInput.getRetryTimeout(); - } - int leftsocketTimeout = loginInput.getSocketTimeout(); int retryCount = 0; @@ -958,16 +953,10 @@ private static SFLoginOutput tokenRequest(SFLoginInput loginInput, TokenRequestT (ArgSupplier) () -> loginInput.getSessionToken() != null ? "******" : null, (ArgSupplier) () -> loginInput.getMasterToken() != null ? "******" : null); - // We want to choose the smaller of the two values between retryTimeout and loginTimeout - int loginRetryTimeout = loginInput.getLoginTimeout(); - if (loginRetryTimeout > loginInput.getRetryTimeout()) { - loginRetryTimeout = loginInput.getRetryTimeout(); - } - String theString = HttpUtil.executeGeneralRequest( postRequest, - loginRetryTimeout, + loginInput.getLoginTimeout(), loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, @@ -1288,16 +1277,10 @@ private static JsonNode federatedFlowStep1(SFLoginInput loginInput) throws Snowf postRequest.addHeader(SF_HEADER_CLIENT_APP_ID, loginInput.getAppId()); postRequest.addHeader(SF_HEADER_CLIENT_APP_VERSION, loginInput.getAppVersion()); - // We want to choose the smaller of the two values between retryTimeout and loginTimeout - int loginRetryTimeout = loginInput.getLoginTimeout(); - if (loginRetryTimeout > loginInput.getRetryTimeout()) { - loginRetryTimeout = loginInput.getRetryTimeout(); - } - final String gsResponse = HttpUtil.executeGeneralRequest( postRequest, - loginRetryTimeout, + loginInput.getLoginTimeout(), loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, @@ -1656,7 +1639,7 @@ public static String generateJWTToken( * @param request the post request * @return true if this is a login/auth request, false otherwise */ - public static boolean isLoginRequest(HttpRequestBase request) { + public static boolean isNewRetryStrategyRequest(HttpRequestBase request) { URI requestURI = request.getURI(); String requestPath = requestURI.getPath(); if (requestPath != null) { diff --git a/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java b/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java index b7194ba0a..b18d1ebcb 100644 --- a/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java +++ b/src/main/java/net/snowflake/client/core/SessionUtilExternalBrowser.java @@ -181,16 +181,10 @@ private String getSSOUrl(int port) throws SFException, SnowflakeSQLException { postRequest.addHeader("accept", "application/json"); - // We want to choose the smaller of the two values between retryTimeout and loginTimeout - int loginRetryTimeout = loginInput.getLoginTimeout(); - if (loginRetryTimeout > loginInput.getRetryTimeout()) { - loginRetryTimeout = loginInput.getRetryTimeout(); - } - String theString = HttpUtil.executeGeneralRequest( postRequest, - loginRetryTimeout, + loginInput.getLoginTimeout(), loginInput.getAuthTimeout(), loginInput.getSocketTimeout(), 0, diff --git a/src/main/java/net/snowflake/client/jdbc/RestRequest.java b/src/main/java/net/snowflake/client/jdbc/RestRequest.java index 3f6b0884a..b31227223 100644 --- a/src/main/java/net/snowflake/client/jdbc/RestRequest.java +++ b/src/main/java/net/snowflake/client/jdbc/RestRequest.java @@ -133,7 +133,7 @@ public static CloseableHttpResponse execute( long startTimePerRequest = startTime; // Used to indicate that this is a login/auth request and will be using the new retry strategy. - boolean isLoginRequest = SessionUtil.isLoginRequest(httpRequest); + boolean isLoginRequest = SessionUtil.isNewRetryStrategyRequest(httpRequest); // total elapsed time due to transient issues. long elapsedMilliForTransientIssues = 0; @@ -426,7 +426,7 @@ public static CloseableHttpResponse execute( (long) backoff.chooseRandom( jitteredBackoffInMilli + backoffInMilli, - Math.pow(retryCount, 2) + jitteredBackoffInMilli); + Math.pow(2, retryCount) + jitteredBackoffInMilli); } else { backoffInMilli = backoff.nextSleepTime(backoffInMilli); } diff --git a/src/test/java/net/snowflake/client/core/SessionUtilTest.java b/src/test/java/net/snowflake/client/core/SessionUtilTest.java index 1846cc1d1..13e7a0a78 100644 --- a/src/test/java/net/snowflake/client/core/SessionUtilTest.java +++ b/src/test/java/net/snowflake/client/core/SessionUtilTest.java @@ -103,7 +103,7 @@ public void testIsLoginRequest() { uriBuilder.setPath(testCase); URI uri = uriBuilder.build(); HttpPost postRequest = new HttpPost(uri); - assertTrue(SessionUtil.isLoginRequest(postRequest)); + assertTrue(SessionUtil.isNewRetryStrategyRequest(postRequest)); } catch (URISyntaxException e) { throw new RuntimeException(e); } @@ -121,7 +121,7 @@ public void testIsLoginRequestInvalidURIPath() { uriBuilder.setPath(testCase); URI uri = uriBuilder.build(); HttpPost postRequest = new HttpPost(uri); - assertFalse(SessionUtil.isLoginRequest(postRequest)); + assertFalse(SessionUtil.isNewRetryStrategyRequest(postRequest)); } catch (URISyntaxException e) { throw new RuntimeException(e); } From 6aaaa0ce2e8ca2196bc1f9227c782ce5abe683ed Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Fri, 3 Nov 2023 15:04:39 -0700 Subject: [PATCH 09/10] review changes --- src/main/java/net/snowflake/client/core/SFSession.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java index bd5bee1e3..7a48f1f11 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -379,7 +379,7 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro case RETRY_TIMEOUT: if (propertyValue != null) { int timeoutValue = (Integer) propertyValue; - if (timeoutValue >= 300) { + if (timeoutValue >= 300 || timeoutValue == 0) { retryTimeout = timeoutValue; } } From 784b1183c574165b17b83d415439e88cc9eec235 Mon Sep 17 00:00:00 2001 From: Jelena Furundzic Date: Fri, 3 Nov 2023 15:29:19 -0700 Subject: [PATCH 10/10] review changes --- src/main/java/net/snowflake/client/core/SFLoginInput.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/src/main/java/net/snowflake/client/core/SFLoginInput.java index d48e3d185..67a4b3aca 100644 --- a/src/main/java/net/snowflake/client/core/SFLoginInput.java +++ b/src/main/java/net/snowflake/client/core/SFLoginInput.java @@ -143,7 +143,7 @@ int getLoginTimeout() { // We want to choose the smaller of the two values between retryTimeout and loginTimeout for the // new retry strategy. SFLoginInput setLoginTimeout(int loginTimeout) { - if (loginTimeout > retryTimeout) { + if (loginTimeout > retryTimeout && retryTimeout != 0) { this.loginTimeout = retryTimeout; } else { this.loginTimeout = loginTimeout;