SNOW-1915351 - retrying after network errors

TestOnly/pom.xml

@@ -397,7 +397,8 @@
             <!-- structured types tests -->
             <exclude>**/structuredtypes/sqldata/*.java</exclude>
             <!-- Latest JDBC driver tests that cannot run with the oldest driver -->
-            <exclude>**/*LatestIT.java</exclude>
+            <exclude>**/*LatestIT.java</exclude>            <!-- Latest JDBC driver tests that cannot run with the oldest driver -->
+            <exclude>**/*WiremockIT.java</exclude>
             <!-- Unit Tests -->
             <exclude>**/*Test.java</exclude>
             <!-- Manual Tests -->

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

@@ -4,6 +4,7 @@
 
 package net.snowflake.client.core;
 
+import static net.snowflake.client.jdbc.ErrorCode.NETWORK_ERROR;
 import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
 import static org.apache.http.client.config.CookieSpecs.DEFAULT;
 import static org.apache.http.client.config.CookieSpecs.IGNORE_COOKIES;
@@ -20,6 +21,7 @@
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.net.Socket;
+import java.net.URISyntaxException;
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
@@ -52,6 +54,7 @@
 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.client.utils.URIBuilder;
 import org.apache.http.config.Registry;
 import org.apache.http.config.RegistryBuilder;
 import org.apache.http.conn.socket.ConnectionSocketFactory;
@@ -73,6 +76,8 @@ public class HttpUtil {
   static final int DEFAULT_MAX_CONNECTIONS_PER_ROUTE = 300;
   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_MALFORMED_RESPONSE_MAX_RETRY_COUNT = 3; // ms
+  static final int DEFAULT_MALFORMED_RESPONSE_RETRY_DELAY = 500; // 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
@@ -81,6 +86,10 @@ public class HttpUtil {
   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";
+  public static final String JDBC_MALFORMED_RESPONSE_MAX_RETRY_COUNT_PROPERTY =
+      "net.snowflake.jdbc.default_malformed_response_max_retry_count";
+  public static final String JDBC_MALFORMED_RESPONSE_RETRY_DELAY =
+      "net.snowflake.jdbc.malformed_response_retry_delay";
 
   private static Duration connectionTimeout;
   private static Duration socketTimeout;
@@ -786,77 +795,121 @@ private static String executeRequestInternal(
     // HttpRequest.toString() contains request URI. Scrub any credentials, if
     // present, before logging
     String requestInfoScrubbed = SecretDetector.maskSASToken(httpRequest.toString());
+    final int malformedResponseMaxRetryCount =
+        SystemUtil.convertSystemPropertyToIntValue(
+            JDBC_MALFORMED_RESPONSE_MAX_RETRY_COUNT_PROPERTY,
+            DEFAULT_MALFORMED_RESPONSE_MAX_RETRY_COUNT);
+    final int malformedResponseRetryDelay =
+        SystemUtil.convertSystemPropertyToIntValue(
+            JDBC_MALFORMED_RESPONSE_RETRY_DELAY, DEFAULT_MALFORMED_RESPONSE_RETRY_DELAY);
+    int retryCount = 0;
+    String responseText = "";
+
+    while (retryCount <= malformedResponseMaxRetryCount) {
+      logger.debug(
+          "Pool: {} Executing: {}",
+          (ArgSupplier) HttpUtil::getHttpClientStats,
+          requestInfoScrubbed);
 
-    logger.debug(
-        "Pool: {} Executing: {}", (ArgSupplier) HttpUtil::getHttpClientStats, requestInfoScrubbed);
-
-    String theString;
-    StringWriter writer = null;
-    CloseableHttpResponse response = null;
-    Stopwatch stopwatch = null;
-
-    if (logger.isDebugEnabled()) {
-      stopwatch = new Stopwatch();
-      stopwatch.start();
-    }
+      StringWriter writer = null;
+      CloseableHttpResponse response = null;
+      Stopwatch stopwatch = null;
 
-    try {
-      response =
-          RestRequest.execute(
-              httpClient,
-              httpRequest,
-              retryTimeout,
-              authTimeout,
-              socketTimeout,
-              maxRetries,
-              injectSocketTimeout,
-              canceling,
-              withoutCookies,
-              includeRetryParameters,
-              includeRequestGuid,
-              retryOnHTTP403,
-              execTimeData);
-      if (logger.isDebugEnabled() && stopwatch != null) {
-        stopwatch.stop();
+      if (logger.isDebugEnabled()) {
+        stopwatch = new Stopwatch();
+        stopwatch.start();
       }
 
-      if (response == null || response.getStatusLine().getStatusCode() != 200) {
-        logger.error("Error executing request: {}", requestInfoScrubbed);
-
-        SnowflakeUtil.logResponseDetails(response, logger);
-
-        if (response != null) {
-          EntityUtils.consume(response.getEntity());
+      try {
+        try {
+          if (includeRetryParameters && retryCount > 0) {
+            URIBuilder uriBuilder = new URIBuilder(httpRequest.getURI());
+            RestRequest.updateRetryParameters(
+                uriBuilder, retryCount, String.valueOf(NETWORK_ERROR), System.currentTimeMillis());
+            httpRequest.setURI(uriBuilder.build());
+          }
+          response =
+              RestRequest.execute(
+                  httpClient,
+                  httpRequest,
+                  retryTimeout,
+                  authTimeout,
+                  socketTimeout,
+                  maxRetries,
+                  injectSocketTimeout,
+                  canceling,
+                  withoutCookies,
+                  includeRetryParameters,
+                  includeRequestGuid,
+                  retryOnHTTP403,
+                  execTimeData);
+          if (logger.isDebugEnabled() && stopwatch != null) {
+            stopwatch.stop();
+          }
+
+          if (response == null || response.getStatusLine().getStatusCode() != 200) {
+            logger.error("Error executing request: {}", requestInfoScrubbed);
+
+            SnowflakeUtil.logResponseDetails(response, logger);
+
+            if (response != null) {
+              EntityUtils.consume(response.getEntity());
+            }
+
+            throw new SnowflakeSQLException(
+                SqlState.IO_ERROR,
+                NETWORK_ERROR.getMessageCode(),
+                "HTTP status="
+                    + ((response != null)
+                        ? response.getStatusLine().getStatusCode()
+                        : "null response"));
+          }
+
+          execTimeData.setResponseIOStreamStart();
+          writer = new StringWriter();
+          try (InputStream ins = response.getEntity().getContent()) {
+            IOUtils.copy(ins, writer, "UTF-8");
+          }
+          responseText = writer.toString();
+          execTimeData.setResponseIOStreamEnd();
+        } catch (URISyntaxException e) {
+          throw new RuntimeException(e);
+        } finally {
+          IOUtils.closeQuietly(writer);
+          IOUtils.closeQuietly(response);
         }
 
-        throw new SnowflakeSQLException(
-            SqlState.IO_ERROR,
-            ErrorCode.NETWORK_ERROR.getMessageCode(),
-            "HTTP status="
-                + ((response != null)
-                    ? response.getStatusLine().getStatusCode()
-                    : "null response"));
-      }
-
-      execTimeData.setResponseIOStreamStart();
-      writer = new StringWriter();
-      try (InputStream ins = response.getEntity().getContent()) {
-        IOUtils.copy(ins, writer, "UTF-8");
+        logger.debug(
+            "Pool: {} Request returned for: {} took {} ms",
+            (ArgSupplier) HttpUtil::getHttpClientStats,
+            requestInfoScrubbed,
+            stopwatch == null ? "n/a" : stopwatch.elapsedMillis());
+
+        break;
+
+      } catch (IOException e) {
+        logger.warn(
+            "Thrown IOException. Message {}. Retry {} . Message {}", e.getMessage(), retryCount);
+        retryCount++;
+        if (retryCount > malformedResponseMaxRetryCount) {
+          String requestIdStr = URLUtil.getRequestIdLogStr(httpRequest.getURI());
+          logger.error(
+              "{}Stop retrying MALFORMED_RESPONSE_ERROR as max retries have been reached for request: {}! Max retry count: {}. Exception message: {}",
+              requestIdStr,
+              requestInfoScrubbed,
+              malformedResponseMaxRetryCount,
+              e.getMessage());
+          throw e;
+        } else {
+          try {
+            Thread.sleep(malformedResponseRetryDelay);
+          } catch (InterruptedException ex) {
+            throw new RuntimeException(ex);
+          }
+        }
       }
-      theString = writer.toString();
-      execTimeData.setResponseIOStreamEnd();
-    } finally {
-      IOUtils.closeQuietly(writer);
-      IOUtils.closeQuietly(response);
     }
-
-    logger.debug(
-        "Pool: {} Request returned for: {} took {} ms",
-        (ArgSupplier) HttpUtil::getHttpClientStats,
-        requestInfoScrubbed,
-        stopwatch == null ? "n/a" : stopwatch.elapsedMillis());
-
-    return theString;
+    return responseText;
   }
 
   // This is a workaround for JDK-7036144.

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

@@ -18,6 +18,7 @@
 import net.snowflake.client.core.HttpUtil;
 import net.snowflake.client.core.SFOCSPException;
 import net.snowflake.client.core.SessionUtil;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.core.URLUtil;
 import net.snowflake.client.core.UUIDUtils;
 import net.snowflake.client.jdbc.telemetryOOB.TelemetryService;
@@ -244,9 +245,7 @@ public static CloseableHttpResponse execute(
           builder.setParameter("target", "htap_simulation");
         }
         if (includeRetryParameters && retryCount > 0) {
-          builder.setParameter("retryCount", String.valueOf(retryCount));
-          builder.setParameter("retryReason", lastStatusCodeForRetry);
-          builder.setParameter("clientStartTime", String.valueOf(startTime));
+          updateRetryParameters(builder, retryCount, lastStatusCodeForRetry, startTime);
         }
 
         // When the auth timeout is set, set the socket timeout as the authTimeout
@@ -628,6 +627,14 @@ public static CloseableHttpResponse execute(
     return response;
   }
 
+  @SnowflakeJdbcInternalApi
+  public static void updateRetryParameters(
+      URIBuilder builder, int retryCount, String lastStatusCodeForRetry, long startTime) {
+    builder.setParameter("retryCount", String.valueOf(retryCount));
+    builder.setParameter("retryReason", lastStatusCodeForRetry);
+    builder.setParameter("clientStartTime", String.valueOf(startTime));
+  }
+
   static long getNewBackoffInMilli(
       long previousBackoffInMilli,
       boolean isLoginRequest,

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

@@ -8,17 +8,23 @@
 import static org.awaitility.Awaitility.await;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
+import com.github.tomakehurst.wiremock.client.WireMock;
 import java.io.File;
+import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.UnsupportedEncodingException;
 import java.net.ServerSocket;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Properties;
 import net.snowflake.client.core.HttpUtil;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
+import org.apache.commons.io.IOUtils;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
@@ -119,6 +125,10 @@ protected void resetWiremock() {
     }
   }
 
+  protected void configureWiremock() {
+    WireMock.configureFor(WIREMOCK_HOST, wiremockHttpPort);
+  }
+
   private static String getWiremockStandAlonePath() {
     return System.getProperty("user.home") + WIREMOCK_M2_PATH;
   }
@@ -237,4 +247,15 @@ protected void addMapping(String mapping) {
       throw new RuntimeException(e);
     }
   }
+
+  protected void importMappingFromResources(String relativePath) {
+    try (InputStream is =
+        new FileInputStream(String.valueOf(Paths.get("src/test/resources/", relativePath)))) {
+      String scenario = IOUtils.toString(Objects.requireNonNull(is), StandardCharsets.UTF_8);
+      importMapping(scenario);
+
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
 }