fix: ensure abort interrupts running queries

aws · Nov 7, 2024 · 93b1467 · 93b1467
1 parent d92b265
commit 93b1467
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java
@@ -25,6 +25,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -69,6 +70,8 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin {
       Collections.unmodifiableSet(new HashSet<String>() {
         {
           addAll(SubscribedMethodHelper.NETWORK_BOUND_METHODS);
+          add("Connection.abort");
+          add("Connection.close");
           add("initHostProvider");
           add("connect");
           add("forceConnect");
@@ -92,7 +95,7 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin {
   protected FailoverMode failoverMode;
   private boolean telemetryFailoverAdditionalTopTraceSetting;
 
-  private boolean closedExplicitly = false;
+  private final AtomicBoolean closedExplicitly = new AtomicBoolean(false);
   protected boolean isClosed = false;
   protected String closedReason = null;
   private final RdsUtils rdsHelper;
@@ -210,7 +213,12 @@ public <T, E extends Exception> T execute(
       final Object[] jdbcMethodArgs)
       throws E {
     if (!this.enableFailoverSetting || canDirectExecute(methodName)) {
-      return jdbcMethodFunc.call();
+      T result = jdbcMethodFunc.call();
+      if (METHOD_ABORT.equals(methodName) || METHOD_CLOSE.equals(methodName)) {
+        this.closedExplicitly.set(true);
+      }
+
+      return result;
     }
 
     if (this.isClosed && !allowedOnClosedConnection(methodName)) {
@@ -228,6 +236,9 @@ public <T, E extends Exception> T execute(
         updateTopology(false);
       }
       result = jdbcMethodFunc.call();
+      if (METHOD_ABORT.equals(methodName) || METHOD_CLOSE.equals(methodName)) {
+        this.closedExplicitly.set(true);
+      }
     } catch (final IllegalStateException e) {
       dealWithIllegalStateException(e, exceptionClass);
     } catch (final Exception e) {
@@ -372,7 +383,7 @@ private void initSettings() {
   }
 
   private void invalidInvocationOnClosedConnection() throws SQLException {
-    if (!this.closedExplicitly) {
+    if (!this.closedExplicitly.get()) {
       this.isClosed = false;
       this.closedReason = null;
       pickNewConnection();
@@ -736,7 +747,7 @@ protected void invalidateCurrentConnection() {
   }
 
   protected void pickNewConnection() throws SQLException {
-    if (this.isClosed && this.closedExplicitly) {
+    if (this.isClosed && this.closedExplicitly.get()) {
       LOGGER.fine(() -> Messages.get("Failover.transactionResolutionUnknownError"));
       return;
     }
@@ -753,6 +764,9 @@ protected void pickNewConnection() throws SQLException {
   }
 
   protected boolean shouldExceptionTriggerConnectionSwitch(final Throwable t) {
+    if (this.closedExplicitly.get()) {
+      return false;
+    }
 
     if (!isFailoverEnabled()) {
       LOGGER.fine(() -> Messages.get("Failover.failoverDisabled"));

diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/AsynchronousMethodsHelper.java b/wrapper/src/main/java/software/amazon/jdbc/util/AsynchronousMethodsHelper.java
@@ -22,6 +22,7 @@
 public class AsynchronousMethodsHelper {
   public static final List<String> ASYNCHRONOUS_METHODS = Arrays.asList(
       "Statement.cancel",
-      "PreparedStatement.cancel"
+      "PreparedStatement.cancel",
+      "Connection.abort"
   );
 }