From 86b0b7f451c23ef5f0bfcf46ebf8beb2f5cfc2d5 Mon Sep 17 00:00:00 2001
From: Dominik Przybysz <dominik.przybysz@snowflake.com>
Date: Fri, 21 Feb 2025 11:49:01 +0100
Subject: [PATCH] SNOW-1943242: Recover from shutdown HTTP connection manager

---
 .../net/snowflake/client/core/HttpUtil.java   | 34 +++++++++++++++++--
 .../client/jdbc/ConnectionLatestIT.java       | 26 ++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/main/java/net/snowflake/client/core/HttpUtil.java b/src/main/java/net/snowflake/client/core/HttpUtil.java
index 85f4a4e70..ebabf8246 100644
--- 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;
@@ -105,9 +106,11 @@ public class HttpUtil {
       new ConcurrentHashMap<>();
 
   /** Handle on the static connection manager, to gather statistics mainly */
+  // TODO SNOW-1943242 consider making it HttpClientSettingsKey bound
   private static PoolingHttpClientConnectionManager connectionManager = null;
 
   /** default request configuration, to be copied on individual requests. */
+  // TODO SNOW-1943242 consider making it HttpClientSettingsKey bound
   private static RequestConfig DefaultRequestConfig = null;
 
   private static boolean socksProxyDisabled = false;
@@ -350,6 +353,8 @@ public static CloseableHttpClient buildHttpClient(
               .build();
 
       // Build a connection manager with enough connections
+      // TODO SNOW-1943242 consider not creating connection manager if already exists and/or
+      //   synchronization
       connectionManager =
           new PoolingHttpClientConnectionManager(
               registry, null, null, null, timeToLive, TimeUnit.SECONDS);
@@ -497,8 +502,33 @@ 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)) {
+      // TODO SNOW-1943242 consider dropping connectionManager
+      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
index 0fd3a7d30..2cefe07d0 100644
--- a/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
+++ b/src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java
@@ -58,6 +58,7 @@
 import net.snowflake.client.category.TestTags;
 import net.snowflake.client.core.HttpClientSettingsKey;
 import net.snowflake.client.core.HttpUtil;
+import net.snowflake.client.core.OCSPMode;
 import net.snowflake.client.core.ObjectMapperFactory;
 import net.snowflake.client.core.QueryStatus;
 import net.snowflake.client.core.SFSession;
@@ -1662,4 +1663,29 @@ 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);
+                }
+              });
+      // next line to override single static HttpUtil.connectionManager
+      HttpUtil.getHttpClient(new HttpClientSettingsKey(OCSPMode.DISABLE_OCSP_CHECKS));
+      try (ResultSet resultSet = statement.executeQuery("Select 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1, resultSet.getInt(1));
+        assertFalse(resultSet.next());
+      }
+    }
+  }
 }