From dc36a0a14bf47e11bba4181ff3f6903c44d0895e Mon Sep 17 00:00:00 2001
From: Janario Oliveira <janario.oliveira@gmail.com>
Date: Thu, 2 Mar 2023 12:51:55 +0100
Subject: [PATCH] Added check with exceptionOverride in case of connection dead

---
 .../java/com/zaxxer/hikari/pool/PoolBase.java | 66 ++++++++++++-------
 1 file changed, 43 insertions(+), 23 deletions(-)

diff --git a/src/main/java/com/zaxxer/hikari/pool/PoolBase.java b/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
index 473997d3c..fa2727b93 100644
--- a/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
+++ b/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
@@ -42,6 +42,7 @@
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static com.zaxxer.hikari.SQLExceptionOverride.Override.DO_NOT_EVICT;
 import static com.zaxxer.hikari.pool.ProxyConnection.*;
 import static com.zaxxer.hikari.util.ClockSource.*;
 import static com.zaxxer.hikari.util.UtilityElf.createInstance;
@@ -145,39 +146,58 @@ void quietlyCloseConnection(final Connection connection, final String closureRea
       }
    }
 
-   boolean isConnectionDead(final Connection connection)
-   {
+   private boolean doIsConnectionDead(final Connection connection) throws SQLException {
+      setNetworkTimeout(connection, validationTimeout);
       try {
-         setNetworkTimeout(connection, validationTimeout);
-         try {
-            final var validationSeconds = (int) Math.max(1000L, validationTimeout) / 1000;
+         final var validationSeconds = (int) Math.max(1000L, validationTimeout) / 1000;
 
-            if (isUseJdbc4Validation) {
-               return !connection.isValid(validationSeconds);
-            }
-
-            try (var statement = connection.createStatement()) {
-               if (isNetworkTimeoutSupported != TRUE) {
-                  setQueryTimeout(statement, validationSeconds);
-               }
+         if (isUseJdbc4Validation) {
+            return !connection.isValid(validationSeconds);
+         }
 
-               statement.execute(config.getConnectionTestQuery());
+         try (var statement = connection.createStatement()) {
+            if (isNetworkTimeoutSupported != TRUE) {
+               setQueryTimeout(statement, validationSeconds);
             }
+
+            statement.execute(config.getConnectionTestQuery());
          }
-         finally {
-            setNetworkTimeout(connection, networkTimeout);
+      }
+      finally {
+         setNetworkTimeout(connection, networkTimeout);
 
-            if (isIsolateInternalQueries && !isAutoCommit) {
-               connection.rollback();
-            }
+         if (isIsolateInternalQueries && !isAutoCommit) {
+            connection.rollback();
          }
+      }
+      return false;
+   }
 
-         return false;
+   private void connectionDeadException(final Connection connection, Exception e) {
+      lastConnectionFailure.set(e);
+      logger.warn("{} - Failed to validate connection {} ({}). Possibly consider using a shorter maxLifetime value.",
+         poolName, connection, e.getMessage());
+   }
+
+   boolean isConnectionDead(final Connection connection)
+   {
+      try {
+         return doIsConnectionDead(connection);
       }
       catch (Exception e) {
-         lastConnectionFailure.set(e);
-         logger.warn("{} - Failed to validate connection {} ({}). Possibly consider using a shorter maxLifetime value.",
-                     poolName, connection, e.getMessage());
+         if (e instanceof SQLException && exceptionOverride != null &&
+            exceptionOverride.adjudicate((SQLException) e) == DO_NOT_EVICT) {
+            // try one more time, in case of failover
+            try {
+               return doIsConnectionDead(connection);
+            }
+            catch (Exception e2) {
+               connectionDeadException(connection, e2);
+               return true;
+            }
+         }
+
+         connectionDeadException(connection, e);
          return true;
       }
    }