fix: statement cancel not working with mysql driver

Author: crystall-bitquill

wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java

@@ -24,9 +24,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
-import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -45,6 +45,7 @@
 import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsPlugin;
 import software.amazon.jdbc.plugin.strategy.fastestresponse.FastestResponseStrategyPlugin;
 import software.amazon.jdbc.profile.ConfigurationProfile;
+import software.amazon.jdbc.util.AsynchronousMethodsHelper;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.SqlMethodAnalyzer;
 import software.amazon.jdbc.util.WrapperUtils;
@@ -91,8 +92,6 @@ public class ConnectionPluginManager implements CanReleaseResources, Wrapper {
   private static final String NOTIFY_CONNECTION_CHANGED_METHOD = "notifyConnectionChanged";
   private static final String NOTIFY_NODE_LIST_CHANGED_METHOD = "notifyNodeListChanged";
   private static final SqlMethodAnalyzer sqlMethodAnalyzer = new SqlMethodAnalyzer();
-  private final ReentrantLock lock = new ReentrantLock();
-
   protected Properties props = new Properties();
   protected List<ConnectionPlugin> plugins;
   protected final @NonNull ConnectionProvider defaultConnProvider;
@@ -148,14 +147,6 @@ public ConnectionPluginManager(
     this.telemetryFactory = telemetryFactory;
   }
 
-  public void lock() {
-    lock.lock();
-  }
-
-  public void unlock() {
-    lock.unlock();
-  }
-
   /**
    * Initialize a chain of {@link ConnectionPlugin} using their corresponding {@link
    * ConnectionPluginFactory}. If {@code PropertyDefinition.PLUGINS} is provided by the user,
@@ -305,13 +296,15 @@ public <T, E extends Exception> T execute(
       final Object[] jdbcMethodArgs)
       throws E {
 
-    final Connection conn = WrapperUtils.getConnectionFromSqlObject(methodInvokeOn);
-    if (conn != null && conn != this.pluginService.getCurrentConnection()
-        && !sqlMethodAnalyzer.isMethodClosingSqlObject(methodName)) {
-      final SQLException e =
-          new SQLException(Messages.get("ConnectionPluginManager.methodInvokedAgainstOldConnection",
-              new Object[] {methodInvokeOn}));
-      throw WrapperUtils.wrapExceptionIfNeeded(exceptionClass, e);
+    // The target driver may block on Statement.getConnection().
+    if (!AsynchronousMethodsHelper.ASYNCHRONOUS_METHODS.contains(methodName)) {
+      final Connection conn = WrapperUtils.getConnectionFromSqlObject(methodInvokeOn);
+      if (!Objects.equals(conn, this.pluginService.getCurrentConnection())
+          && !sqlMethodAnalyzer.isMethodClosingSqlObject(methodName)) {
+        throw WrapperUtils.wrapExceptionIfNeeded(
+            exceptionClass,
+            new SQLException(Messages.get("ConnectionPluginManager.invokedAgainstOldConnection", new Object[] {methodInvokeOn})));
+      }
     }
 
     return executeWithSubscribedPlugins(

wrapper/src/main/java/software/amazon/jdbc/util/AsynchronousMethodsHelper.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package software.amazon.jdbc.util;
+
+import java.util.Collections;
+import java.util.List;
+
+public class AsynchronousMethodsHelper {
+  public static final List<String> ASYNCHRONOUS_METHODS = Collections.singletonList(
+      "Statement.cancel"
+  );
+}

wrapper/src/main/java/software/amazon/jdbc/util/WrapperUtils.java

@@ -50,6 +50,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.locks.ReentrantLock;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.ConnectionPluginManager;
 import software.amazon.jdbc.JdbcCallable;
@@ -84,6 +85,8 @@ public class WrapperUtils {
   private static final ConcurrentMap<Class<?>, Boolean> isJdbcInterfaceCache =
       new ConcurrentHashMap<>();
 
+  private static final ReentrantLock lock = new ReentrantLock();
+
   private static final Map<Class<?>, Class<?>> availableWrappers =
       new HashMap<Class<?>, Class<?>>() {
         {
@@ -182,7 +185,9 @@ public static <T> T executeWithPlugins(
       final JdbcCallable<T, RuntimeException> jdbcMethodFunc,
       final Object... jdbcMethodArgs) {
 
-    pluginManager.lock();
+    if (!AsynchronousMethodsHelper.ASYNCHRONOUS_METHODS.contains(methodName)) {
+      lock.lock();
+    }
     TelemetryFactory telemetryFactory = pluginManager.getTelemetryFactory();
     TelemetryContext context = null;
 
@@ -208,7 +213,9 @@ public static <T> T executeWithPlugins(
         throw new RuntimeException(e);
       }
     } finally {
-      pluginManager.unlock();
+      if (lock.isHeldByCurrentThread()) {
+        lock.unlock();
+      }
       if (context != null) {
         context.closeContext();
       }
@@ -225,7 +232,9 @@ public static <T, E extends Exception> T executeWithPlugins(
       final Object... jdbcMethodArgs)
       throws E {
 
-    pluginManager.lock();
+    if (!AsynchronousMethodsHelper.ASYNCHRONOUS_METHODS.contains(methodName)) {
+      lock.lock();
+    }
     TelemetryFactory telemetryFactory = pluginManager.getTelemetryFactory();
     TelemetryContext context = null;
 
@@ -251,7 +260,9 @@ public static <T, E extends Exception> T executeWithPlugins(
       }
 
     } finally {
-      pluginManager.unlock();
+      if (lock.isHeldByCurrentThread()) {
+        lock.unlock();
+      }
       if (context != null) {
         context.closeContext();
       }

wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties

@@ -90,7 +90,7 @@ ConnectionStringHostListProvider.errorIdentifyConnection=An error occurred while
 ConnectionPluginManager.releaseResources=Releasing resources.
 ConnectionPluginManager.unknownPluginCode=Unknown plugin code: ''{0}''.
 ConnectionPluginManager.unableToLoadPlugin=Unable to load connection plugin factory: ''{0}''.
-ConnectionPluginManager.methodInvokedAgainstOldConnection=The internal connection has changed since ''{0}'' was created. This is likely due to failover or read-write splitting functionality. To ensure you are using the updated connection, please re-create Statement and ResultSet objects after failover and/or calling setReadOnly.
+ConnectionPluginManager.invokedAgainstOldConnection=The internal connection has changed since ''{0}'' was created. This is likely due to failover or read-write splitting functionality. To ensure you are using the updated connection, please re-create Statement and ResultSet objects after failover and/or calling setReadOnly.
 
 # Connection Provider
 ConnectionProvider.noConnection=The target driver did not return a connection.

wrapper/src/test/java/software/amazon/jdbc/util/WrapperUtilsTest.java

@@ -20,6 +20,7 @@
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
@@ -54,19 +55,9 @@ public class WrapperUtilsTest {
   @BeforeEach
   @SuppressWarnings("unchecked")
   void init() {
-    final ReentrantLock pluginManagerLock = new ReentrantLock();
     final ReentrantLock testLock = new ReentrantLock();
     closeable = MockitoAnnotations.openMocks(this);
 
-    doAnswer(invocation -> {
-      pluginManagerLock.lock();
-      return null;
-    }).when(pluginManager).lock();
-    doAnswer(invocation -> {
-      pluginManagerLock.unlock();
-      return null;
-    }).when(pluginManager).unlock();
-
     doAnswer(invocation -> {
       boolean lockIsFree = testLock.tryLock();
       if (!lockIsFree) {
@@ -79,7 +70,7 @@ void init() {
         any(Class.class),
         any(Class.class),
         any(Object.class),
-        any(String.class),
+        argThat(methodName -> !AsynchronousMethodsHelper.ASYNCHRONOUS_METHODS.contains(methodName)),
         any(JdbcCallable.class),
         any(Object[].class));
 
@@ -93,23 +84,39 @@ void cleanUp() throws Exception {
     closeable.close();
   }
 
+  Integer callCancelExecuteWithPlugins() {
+    return callExecuteWithPlugins("Statement.cancel");
+  }
+
   Integer callExecuteWithPlugins() {
+    return callExecuteWithPlugins("methodName");
+  }
+
+  Integer callExecuteWithPlugins(String methodName) {
     return WrapperUtils.executeWithPlugins(
         Integer.class,
         pluginManager,
         object,
-        "methodName",
+        methodName,
         () -> 1);
   }
 
+  Integer callCancelExecuteWithPluginsWithException() {
+    return callExecuteWithPluginsWithException("Statement.cancel");
+  }
+
   Integer callExecuteWithPluginsWithException() {
+    return callExecuteWithPluginsWithException("methodName");
+  }
+
+  Integer callExecuteWithPluginsWithException(String methodName) {
     try {
       return WrapperUtils.executeWithPlugins(
           Integer.class,
           SQLException.class,
           pluginManager,
           object,
-          "methodName",
+          methodName,
           () -> 1);
     } catch (SQLException e) {
       fail();
@@ -144,6 +151,30 @@ void testExecutesWithPluginsWithExceptionIsSequential() {
     }
   }
 
+  @Test
+  void testCancelStatementIsNotBlockedExecute() {
+    List<CompletableFuture<Integer>> futures = new ArrayList<>();
+
+    futures.add(CompletableFuture.supplyAsync(this::callExecuteWithPlugins));
+    futures.add(CompletableFuture.supplyAsync(this::callCancelExecuteWithPlugins));
+
+    for (CompletableFuture<Integer> future : futures) {
+      future.join();
+    }
+  }
+
+  @Test
+  void testCancelStatementIsNotBlockedExecuteWithException() {
+    List<CompletableFuture<Integer>> futures = new ArrayList<>();
+
+    futures.add(CompletableFuture.supplyAsync(this::callExecuteWithPluginsWithException));
+    futures.add(CompletableFuture.supplyAsync(this::callCancelExecuteWithPluginsWithException));
+
+    for (CompletableFuture<Integer> future : futures) {
+      future.join();
+    }
+  }
+
   @Test
   void getConnectionFromSqlObjectChecksStatementNotClosed() throws Exception {
     final Statement mockClosedStatement = mock(Statement.class);