SNOW-1943242 JDBC Caches CloseableHttpClient Indefinitely After Connection Pool Shut Down Error (#2224)

Co-authored-by: Harry Xi <[email protected]>

Author: sfc-gh-ext-simba-vb

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

@@ -46,7 +46,6 @@
 import org.apache.http.auth.UsernamePasswordCredentials;
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.client.config.RequestConfig;
-import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpRequestBase;
 import org.apache.http.config.Registry;
 import org.apache.http.config.RegistryBuilder;
@@ -745,7 +744,10 @@ static String executeRequestWithoutCookies(
         getHttpClient(ocspAndProxyKey, null),
         new ExecTimeTelemetryData(),
         null,
-        sfSession);
+        sfSession,
+        ocspAndProxyKey,
+        null,
+        false);
   }
 
   /**
@@ -842,7 +844,10 @@ public static String executeGeneralRequestOmitRequestGuid(
         getHttpClient(ocspAndProxyAndGzipKey, null),
         new ExecTimeTelemetryData(),
         null,
-        sfSession);
+        sfSession,
+        ocspAndProxyAndGzipKey,
+        null,
+        false);
   }
 
   /**
@@ -957,7 +962,10 @@ public static String executeGeneralRequest(
         httpClient,
         new ExecTimeTelemetryData(),
         null,
-        sfSession);
+        sfSession,
+        null,
+        null,
+        false);
   }
 
   /**
@@ -1161,7 +1169,10 @@ public static String executeRequest(
         getHttpClient(ocspAndProxyKey, null),
         execTimeData,
         retryContextManager,
-        sfSession);
+        sfSession,
+        ocspAndProxyKey,
+        null,
+        false);
   }
 
   /**
@@ -1185,6 +1196,10 @@ public static String executeRequest(
    * @param httpClient client object used to communicate with other machine
    * @param retryContextManager RetryContext used to customize retry handling functionality
    * @param sfSession the session associated with the request
+   * @param key HttpClientSettingsKey object
+   * @param httpHeaderCustomizer HttpHeadersCustomizer object for customization of HTTP headers for
+   *     requests sent by the Snowflake JDBC driver.
+   * @param isHttpClientWithoutDecompression flag for create client without Decompression
    * @return response in String
    * @throws SnowflakeSQLException if Snowflake error occurs
    * @throws IOException raises if a general IO error occurs
@@ -1204,15 +1219,17 @@ private static String executeRequestInternal(
       CloseableHttpClient httpClient,
       ExecTimeTelemetryData execTimeData,
       RetryContextManager retryContextManager,
-      SFBaseSession sfSession)
+      SFBaseSession sfSession,
+      HttpClientSettingsKey key,
+      List<HttpHeadersCustomizer> httpHeaderCustomizer,
+      boolean isHttpClientWithoutDecompression)
       throws SnowflakeSQLException, IOException {
     String requestInfoScrubbed = SecretDetector.maskSASToken(httpRequest.toString());
     String responseText = "";
 
     logger.debug(
         "Pool: {} Executing: {}", (ArgSupplier) HttpUtil::getHttpClientStats, requestInfoScrubbed);
 
-    CloseableHttpResponse response = null;
     Stopwatch stopwatch = null;
 
     String requestIdStr = URLUtil.getRequestIdLogStr(httpRequest.getURI());
@@ -1235,7 +1252,14 @@ private static String executeRequestInternal(
             .build();
     responseText =
         RestRequest.executeWithRetries(
-                httpClient, httpRequest, context, execTimeData, retryContextManager)
+                httpClient,
+                httpRequest,
+                context,
+                execTimeData,
+                retryContextManager,
+                key,
+                httpHeaderCustomizer,
+                isHttpClientWithoutDecompression)
             .getUnpackedCloseableHttpResponse();
 
     logger.debug(

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

@@ -148,7 +148,10 @@ else if (context.getQrmk() != null) {
                 true, // no retry on http request
                 false,
                 new ExecTimeTelemetryData(),
-                session)
+                session,
+                context.getChunkDownloader().getHttpClientSettingsKey(),
+                headersCustomizers,
+                false)
             .getHttpResponse();
 
     logger.debug(

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

@@ -24,6 +24,7 @@
 import net.snowflake.client.core.Event;
 import net.snowflake.client.core.EventUtil;
 import net.snowflake.client.core.ExecTimeTelemetryData;
+import net.snowflake.client.core.HttpClientSettingsKey;
 import net.snowflake.client.core.HttpExecutingContext;
 import net.snowflake.client.core.HttpExecutingContextBuilder;
 import net.snowflake.client.core.HttpResponseContextDto;
@@ -197,7 +198,10 @@ public static CloseableHttpResponse execute(
         false, // noRetry
         execTimeTelemetryData,
         null,
-        sfSession);
+        sfSession,
+        null,
+        null,
+        false);
   }
 
   /**
@@ -312,7 +316,10 @@ public static CloseableHttpResponse execute(
         noRetry,
         execTimeData,
         null,
-        sfSession);
+        sfSession,
+        null,
+        null,
+        false);
   }
 
   /**
@@ -427,7 +434,10 @@ public static CloseableHttpResponse execute(
         false, // noRetry
         execTimeData,
         retryContextManager,
-        sfSession);
+        sfSession,
+        null,
+        null,
+        false);
   }
 
   /**
@@ -488,9 +498,11 @@ public static CloseableHttpResponse execute(
         noRetry,
         execTimeData,
         retryManager,
-        null);
+        null,
+        null,
+        null,
+        false);
   }
-
   /**
    * Execute an HTTP request with retry logic.
    *
@@ -512,6 +524,10 @@ public static CloseableHttpResponse execute(
    * @param retryManager RetryContextManager - object allowing to optionally pass custom logic that
    *     should be executed before and/or after the retry
    * @param sfSession the session associated with the request
+   * @param key HttpClientSettingsKey object
+   * @param httpHeaderCustomizer HttpHeadersCustomizer object for customization of HTTP headers for
+   *     requests sent by the Snowflake JDBC driver.
+   * @param isHttpClientWithoutDecompression flag for create client without Decompression
    * @return HttpResponse Object get from server
    * @throws net.snowflake.client.jdbc.SnowflakeSQLException Request timeout Exception or Illegal
    *     State Exception i.e. connection is already shutdown etc
@@ -533,7 +549,10 @@ public static CloseableHttpResponse execute(
       boolean noRetry,
       ExecTimeTelemetryData execTimeData,
       RetryContextManager retryManager,
-      SFBaseSession sfSession)
+      SFBaseSession sfSession,
+      HttpClientSettingsKey key,
+      List<HttpHeadersCustomizer> httpHeaderCustomizer,
+      boolean isHttpClientWithoutDecompression)
       throws SnowflakeSQLException {
     return executeWithRetries(
             httpClient,
@@ -550,7 +569,10 @@ public static CloseableHttpResponse execute(
             retryHTTP403, // retry on HTTP 403
             noRetry,
             new ExecTimeTelemetryData(),
-            sfSession)
+            sfSession,
+            key,
+            httpHeaderCustomizer,
+            isHttpClientWithoutDecompression)
         .getHttpResponse();
   }
 
@@ -738,7 +760,10 @@ public static HttpResponseContextDto executeWithRetries(
       boolean retryHTTP403,
       boolean unpackResponse,
       ExecTimeTelemetryData execTimeTelemetryData,
-      SFBaseSession sfSession)
+      SFBaseSession sfSession,
+      HttpClientSettingsKey key,
+      List<HttpHeadersCustomizer> httpHeaderCustomizer,
+      boolean isHttpClientWithoutDecompression)
       throws SnowflakeSQLException {
     return executeWithRetries(
         httpClient,
@@ -756,7 +781,10 @@ public static HttpResponseContextDto executeWithRetries(
         false,
         unpackResponse,
         execTimeTelemetryData,
-        sfSession);
+        sfSession,
+        key,
+        httpHeaderCustomizer,
+        isHttpClientWithoutDecompression);
   }
 
   /**
@@ -799,7 +827,10 @@ public static HttpResponseContextDto executeWithRetries(
       boolean noRetry,
       boolean unpackResponse,
       ExecTimeTelemetryData execTimeTelemetryData,
-      SFBaseSession sfSession)
+      SFBaseSession sfSession,
+      HttpClientSettingsKey key,
+      List<HttpHeadersCustomizer> httpHeaderCustomizer,
+      boolean isHttpClientWithoutDecompression)
       throws SnowflakeSQLException {
     String requestIdStr = URLUtil.getRequestIdLogStr(httpRequest.getURI());
     String requestInfoScrubbed = SecretDetector.maskSASToken(httpRequest.toString());
@@ -820,7 +851,15 @@ public static HttpResponseContextDto executeWithRetries(
             .loginRequest(SessionUtil.isNewRetryStrategyRequest(httpRequest))
             .withSfSession(sfSession)
             .build();
-    return executeWithRetries(httpClient, httpRequest, context, execTimeTelemetryData, null);
+    return executeWithRetries(
+        httpClient,
+        httpRequest,
+        context,
+        execTimeTelemetryData,
+        null,
+        key,
+        httpHeaderCustomizer,
+        isHttpClientWithoutDecompression);
   }
 
   /**
@@ -840,7 +879,10 @@ public static HttpResponseContextDto executeWithRetries(
       HttpRequestBase httpRequest,
       HttpExecutingContext httpExecutingContext,
       ExecTimeTelemetryData execTimeData,
-      RetryContextManager retryManager)
+      RetryContextManager retryManager,
+      HttpClientSettingsKey key,
+      List<HttpHeadersCustomizer> httpHeaderCustomizer,
+      boolean isHttpClientWithoutDecompression)
       throws SnowflakeSQLException {
     Stopwatch networkComunnicationStapwatch = null;
     Stopwatch requestReponseStopWatch = null;
@@ -908,7 +950,26 @@ public static HttpResponseContextDto executeWithRetries(
         responseDto.setHttpResponse(response);
         execTimeData.setHttpClientEnd();
       } catch (Exception ex) {
-        responseDto.setSavedEx(handlingNotRetryableException(ex, httpExecutingContext));
+        if (ex instanceof IllegalStateException) {
+          // if exception is caused by illegal state, e.g shutdown of http client
+          // because of closing of connection, then recreate the http client and remove existing
+          // from the cache.
+          logger.warn(
+              "IllegalStateException encountered while processing the HTTP request."
+                  + " The HttpClient was shut down due to connection closure. "
+                  + "Attempting to rebuild the HttpClient and retry the request.");
+          // Clear the httpClient cache.
+          HttpUtil.httpClient.remove(key);
+          // rebuild the http client.
+          if (isHttpClientWithoutDecompression) {
+            httpClient = HttpUtil.getHttpClientWithoutDecompression(key, httpHeaderCustomizer);
+          } else {
+            httpClient = HttpUtil.getHttpClient(key, httpHeaderCustomizer);
+          }
+          continue;
+        } else {
+          responseDto.setSavedEx(handlingNotRetryableException(ex, httpExecutingContext));
+        }
       } finally {
         // Reset the socket timeout to its original value if it is not the
         // very first iteration.

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

@@ -254,7 +254,10 @@ public void download(
                   true, // retry on HTTP 403
                   false,
                   new ExecTimeTelemetryData(),
-                  session);
+                  session,
+                  session.getHttpClientKey(),
+                  session.getHttpHeadersCustomizers(),
+                  true);
           HttpResponse response = responseDto.getHttpResponse();
 
           logger.debug(
@@ -446,7 +449,10 @@ public InputStream downloadToStream(
                       true, // retry on HTTP 403
                       false,
                       new ExecTimeTelemetryData(),
-                      session)
+                      session,
+                      session.getHttpClientKey(),
+                      session.getHttpHeadersCustomizers(),
+                      true)
                   .getHttpResponse();
 
           logger.debug(
@@ -943,7 +949,10 @@ private void uploadWithPresignedUrl(
                   true, // retry on HTTP 403
                   true, // disable retry
                   new ExecTimeTelemetryData(),
-                  session)
+                  session,
+                  ocspAndProxyKey,
+                  session.getHttpHeadersCustomizers(),
+                  false)
               .getHttpResponse();
 
       logger.debug(

src/test/java/net/snowflake/client/jdbc/ChunkDownloaderS3RetryUrlLatestIT.java

@@ -83,7 +83,10 @@ public void testParamsInRetryS3Url() throws Exception {
         true,
         false,
         new ExecTimeTelemetryData(),
-        null); // retry HTTP 403
+        null, // retry HTTP 403
+        sfContext.getChunkDownloader().getHttpClientSettingsKey(),
+        null,
+        false);
 
     assertFalse(getRequest.containsHeader("retryCount"));
     assertFalse(getRequest.containsHeader("retryReason"));