Snow 1853185 - Okta http retry storm - retries implementation for Okta auth (#2064)

Author: sfc-gh-fpawlowski

TestOnly/pom.xml

@@ -22,6 +22,7 @@
     <surefire.version>3.5.1</surefire.version>
     <mockito.version>3.5.6</mockito.version>
     <netty.version>4.1.118.Final</netty.version>
+    <apache.httpclient.version>4.5.14</apache.httpclient.version>
     <shadeBase>net.snowflake.client.jdbc.internal</shadeBase>
   </properties>
 
@@ -31,6 +32,11 @@
       <artifactId>netty-common</artifactId>
       <version>${netty.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.httpcomponents</groupId>
+      <artifactId>httpclient</artifactId>
+      <version>${apache.httpclient.version}</version>
+    </dependency>
     <dependency>
       <groupId>io.netty</groupId>
       <artifactId>netty-buffer</artifactId>

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

@@ -28,6 +28,7 @@
 import javax.net.ssl.TrustManager;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.RestRequest;
+import net.snowflake.client.jdbc.RetryContextManager;
 import net.snowflake.client.jdbc.SnowflakeDriver;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
 import net.snowflake.client.jdbc.SnowflakeUtil;
@@ -582,7 +583,8 @@ public static boolean isSocksProxyDisabled() {
    * @param retryTimeout retry timeout
    * @param authTimeout authenticator specific timeout
    * @param socketTimeout socket timeout (in ms)
-   * @param retryCount retry count for the request
+   * @param retryCount max retry count for the request - if it is set to 0, it will be ignored and
+   *     only retryTimeout will determine when to end the retries
    * @param injectSocketTimeout injecting socket timeout
    * @param canceling canceling?
    * @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
@@ -614,7 +616,8 @@ static String executeRequestWithoutCookies(
         true, // guid? (do we need this?)
         false, // no retry on HTTP 403
         getHttpClient(ocspAndProxyKey),
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        null);
   }
 
   /**
@@ -624,7 +627,8 @@ static String executeRequestWithoutCookies(
    * @param retryTimeout retry timeout
    * @param authTimeout authenticator specific timeout
    * @param socketTimeout socket timeout (in ms)
-   * @param retryCount retry count for the request
+   * @param retryCount max retry count for the request - if it is set to 0, it will be ignored and
+   *     only retryTimeout will determine when to end the retries
    * @param ocspAndProxyAndGzipKey OCSP mode and proxy settings for httpclient
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
@@ -638,6 +642,41 @@ public static String executeGeneralRequest(
       int retryCount,
       HttpClientSettingsKey ocspAndProxyAndGzipKey)
       throws SnowflakeSQLException, IOException {
+    return executeGeneralRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        retryCount,
+        ocspAndProxyAndGzipKey,
+        null);
+  }
+
+  /**
+   * Executes an HTTP request for Snowflake.
+   *
+   * @param httpRequest HttpRequestBase
+   * @param retryTimeout retry timeout
+   * @param authTimeout authenticator specific timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param retryCount max retry count for the request - if it is set to 0, it will be ignored and
+   *     only retryTimeout will determine when to end the retries
+   * @param ocspAndProxyAndGzipKey OCSP mode and proxy settings for httpclient
+   * @param retryContextManager RetryContext used to customize retry handling functionality
+   * @return response
+   * @throws SnowflakeSQLException if Snowflake error occurs
+   * @throws IOException raises if a general IO error occurs
+   */
+  @SnowflakeJdbcInternalApi
+  public static String executeGeneralRequest(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int retryCount,
+      HttpClientSettingsKey ocspAndProxyAndGzipKey,
+      RetryContextManager retryContextManager)
+      throws SnowflakeSQLException, IOException {
     logger.debug("Executing general request");
     return executeRequest(
         httpRequest,
@@ -650,7 +689,8 @@ public static String executeGeneralRequest(
         false, // no retry parameter
         false, // no retry on HTTP 403
         ocspAndProxyAndGzipKey,
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        retryContextManager);
   }
 
   /**
@@ -660,7 +700,8 @@ public static String executeGeneralRequest(
    * @param retryTimeout retry timeout
    * @param authTimeout authenticator specific timeout
    * @param socketTimeout socket timeout (in ms)
-   * @param retryCount retry count for the request
+   * @param retryCount max retry count for the request - if it is set to 0, it will be ignored and
+   *     only retryTimeout will determine when to end the retries
    * @param httpClient client object used to communicate with other machine
    * @return response
    * @throws SnowflakeSQLException if Snowflake error occurs
@@ -688,7 +729,8 @@ public static String executeGeneralRequest(
         true, // include request GUID
         false, // no retry on HTTP 403
         httpClient,
-        new ExecTimeTelemetryData());
+        new ExecTimeTelemetryData(),
+        null);
   }
 
   /**
@@ -722,6 +764,54 @@ public static String executeRequest(
       HttpClientSettingsKey ocspAndProxyKey,
       ExecTimeTelemetryData execTimeData)
       throws SnowflakeSQLException, IOException {
+    return executeRequest(
+        httpRequest,
+        retryTimeout,
+        authTimeout,
+        socketTimeout,
+        maxRetries,
+        injectSocketTimeout,
+        canceling,
+        includeRetryParameters,
+        retryOnHTTP403,
+        ocspAndProxyKey,
+        execTimeData,
+        null);
+  }
+
+  /**
+   * Executes an HTTP request for Snowflake.
+   *
+   * @param httpRequest HttpRequestBase
+   * @param retryTimeout retry timeout
+   * @param authTimeout authenticator timeout
+   * @param socketTimeout socket timeout (in ms)
+   * @param maxRetries retry count for the request
+   * @param injectSocketTimeout injecting socket timeout
+   * @param canceling canceling?
+   * @param includeRetryParameters whether to include retry parameters in retried requests
+   * @param retryOnHTTP403 whether to retry on HTTP 403 or not
+   * @param ocspAndProxyKey OCSP mode and proxy settings for httpclient
+   * @param execTimeData query execution time telemetry data object
+   * @param retryContextManager RetryContext used to customize retry handling functionality
+   * @return response
+   * @throws SnowflakeSQLException if Snowflake error occurs
+   * @throws IOException raises if a general IO error occurs
+   */
+  public static String executeRequest(
+      HttpRequestBase httpRequest,
+      int retryTimeout,
+      int authTimeout,
+      int socketTimeout,
+      int maxRetries,
+      int injectSocketTimeout,
+      AtomicBoolean canceling,
+      boolean includeRetryParameters,
+      boolean retryOnHTTP403,
+      HttpClientSettingsKey ocspAndProxyKey,
+      ExecTimeTelemetryData execTimeData,
+      RetryContextManager retryContextManager)
+      throws SnowflakeSQLException, IOException {
     boolean ocspEnabled = !(ocspAndProxyKey.getOcspMode().equals(OCSPMode.DISABLE_OCSP_CHECKS));
     logger.debug("Executing request with OCSP enabled: {}", ocspEnabled);
     execTimeData.setOCSPStatus(ocspEnabled);
@@ -738,7 +828,8 @@ public static String executeRequest(
         true, // include request GUID
         retryOnHTTP403,
         getHttpClient(ocspAndProxyKey),
-        execTimeData);
+        execTimeData,
+        retryContextManager);
   }
 
   /**
@@ -760,6 +851,7 @@ public static String executeRequest(
    * @param includeRequestGuid whether to include request_guid
    * @param retryOnHTTP403 whether to retry on HTTP 403
    * @param httpClient client object used to communicate with other machine
+   * @param retryContextManager RetryContext used to customize retry handling functionality
    * @return response in String
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -777,7 +869,8 @@ private static String executeRequestInternal(
       boolean includeRequestGuid,
       boolean retryOnHTTP403,
       CloseableHttpClient httpClient,
-      ExecTimeTelemetryData execTimeData)
+      ExecTimeTelemetryData execTimeData,
+      RetryContextManager retryContextManager)
       throws SnowflakeSQLException, IOException {
     // HttpRequest.toString() contains request URI. Scrub any credentials, if
     // present, before logging
@@ -811,7 +904,8 @@ private static String executeRequestInternal(
               includeRetryParameters,
               includeRequestGuid,
               retryOnHTTP403,
-              execTimeData);
+              execTimeData,
+              retryContextManager);
       if (logger.isDebugEnabled() && stopwatch != null) {
         stopwatch.stop();
       }
@@ -825,6 +919,7 @@ private static String executeRequestInternal(
           EntityUtils.consume(response.getEntity());
         }
 
+        // We throw here exception if timeout was reached for login
         throw new SnowflakeSQLException(
             SqlState.IO_ERROR,
             ErrorCode.NETWORK_ERROR.getMessageCode(),