SNOW-1943242: Recover from shutdown HTTP connection manager

snowflakedb · Feb 21, 2025 · 00bb57d · 00bb57d
1 parent dcb60b2
commit 00bb57d
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/src/main/java/net/snowflake/client/core/HttpUtil.java b/src/main/java/net/snowflake/client/core/HttpUtil.java
@@ -17,6 +17,7 @@
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.net.Socket;
@@ -497,8 +498,32 @@ public static CloseableHttpClient initHttpClientWithoutDecompression(
    */
   public static CloseableHttpClient initHttpClient(HttpClientSettingsKey key, File ocspCacheFile) {
     updateRoutePlanner(key);
-    return httpClient.computeIfAbsent(
-        key, k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled()));
+    CloseableHttpClient closeableHttpClient =
+        httpClient.computeIfAbsent(
+            key, k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled()));
+    if (detectClosedConnectionManager(closeableHttpClient)) {
+      httpClient.remove(key);
+      return httpClient.computeIfAbsent(
+          key, k -> buildHttpClient(key, ocspCacheFile, key.getGzipDisabled()));
+    }
+    return closeableHttpClient;
+  }
+
+  private static boolean detectClosedConnectionManager(CloseableHttpClient closeableHttpClient)
+      throws RuntimeException {
+    // There is no simple way to detect is the connection manager is closed...
+    try {
+      Field connManagerField = closeableHttpClient.getClass().getDeclaredField("connManager");
+      connManagerField.setAccessible(true);
+      PoolingHttpClientConnectionManager connectionManager =
+          (PoolingHttpClientConnectionManager) connManagerField.get(closeableHttpClient);
+      Field isShutdownField = connectionManager.getClass().getDeclaredField("isShutDown");
+      isShutdownField.setAccessible(true);
+      AtomicBoolean isShutdown = (AtomicBoolean) isShutdownField.get(connectionManager);
+      return isShutdown.get();
+    } catch (NoSuchFieldException | IllegalAccessException 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
@@ -1662,4 +1662,27 @@ public void testDisableOCSPChecksMode() throws SQLException {
     assertThat(
         thrown.getErrorCode(), anyOf(is(INVALID_CONNECTION_INFO_CODE), is(BAD_REQUEST_GS_CODE)));
   }
+
+  /** Added after version 3.22.0 */
+  @Test
+  public void testRecoverFromClosedHttpConnectionManager() throws SQLException {
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
+      HttpUtil.httpClient
+          .values()
+          .forEach(
+              closeableHttpClient -> {
+                try {
+                  closeableHttpClient.close();
+                } catch (IOException e) {
+                  throw new RuntimeException(e);
+                }
+              });
+      try (ResultSet resultSet = statement.executeQuery("Select 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1, resultSet.getInt(1));
+        assertFalse(resultSet.next());
+      }
+    }
+  }
 }