Use try with resources to simplify locking code

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

@@ -25,7 +25,6 @@
 import java.util.Map;
 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;
@@ -46,6 +45,7 @@
 import software.amazon.jdbc.util.SqlMethodAnalyzer;
 import software.amazon.jdbc.util.SqlState;
 import software.amazon.jdbc.util.StringUtils;
+import software.amazon.jdbc.util.ResourceLock;
 import software.amazon.jdbc.util.WrapperUtils;
 import software.amazon.jdbc.wrapper.ConnectionWrapper;
 
@@ -82,7 +82,7 @@ public class ConnectionPluginManager implements CanReleaseResources {
   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();
+  private final ResourceLock lock = new ResourceLock();
 
   protected Properties props = new Properties();
   protected ArrayList<ConnectionPlugin> plugins;
@@ -125,12 +125,15 @@ public ConnectionPluginManager(ConnectionProvider connectionProvider, Connection
     this.connectionWrapper = connectionWrapper;
   }
 
-  public void lock() {
-    lock.lock();
+  public ResourceLock acquireLock() {
+    return lock.obtain();
   }
 
-  public void unlock() {
-    lock.unlock();
+  /*
+  For testing only
+   */
+  public void releaseLock() {
+    lock.close();
   }
 
   /**

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

@@ -0,0 +1,54 @@
+/*
+ * 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.
+ */
+/*
+ * portions Copyright (c) 2022, PostgreSQL Global Development Group
+ * See the LICENSE file in the project root for more information.
+ */
+
+package software.amazon.jdbc.util;
+
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Extends a ReentrantLock for use in try-with-resources block.
+ *
+ * <h2>Example use</h2>
+ * <pre>{@code
+ *
+ *   try (ResourceLock ignore = lock.obtain()) {
+ *     // do something while holding the resource lock
+ *   }
+ *
+ * }</pre>
+ */
+public final class ResourceLock extends ReentrantLock implements AutoCloseable {
+
+  /**
+   * Obtain a lock and return the ResourceLock for use in try-with-resources block.
+   */
+  public ResourceLock obtain() {
+    lock();
+    return this;
+  }
+
+  /**
+   * Unlock on exit of try-with-resources block.
+   */
+  @Override
+  public void close() {
+    this.unlock();
+  }
+}

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

@@ -180,9 +180,8 @@ public static <T> T executeWithPlugins(
       final JdbcCallable<T, RuntimeException> jdbcMethodFunc,
       final Object... jdbcMethodArgs) {
 
-    pluginManager.lock();
 
-    try {
+    try (ResourceLock ignore = pluginManager.acquireLock()) {
       final Object[] argsCopy =
           jdbcMethodArgs == null ? null : Arrays.copyOf(jdbcMethodArgs, jdbcMethodArgs.length);
 
@@ -200,9 +199,6 @@ public static <T> T executeWithPlugins(
       } catch (final InstantiationException e) {
         throw new RuntimeException(e);
       }
-
-    } finally {
-      pluginManager.unlock();
     }
   }
 
@@ -216,9 +212,7 @@ public static <T, E extends Exception> T executeWithPlugins(
       final Object... jdbcMethodArgs)
       throws E {
 
-    pluginManager.lock();
-
-    try {
+    try (ResourceLock lock = pluginManager.acquireLock()){
       final Object[] argsCopy =
           jdbcMethodArgs == null ? null : Arrays.copyOf(jdbcMethodArgs, jdbcMethodArgs.length);
 
@@ -232,8 +226,6 @@ public static <T, E extends Exception> T executeWithPlugins(
         throw new RuntimeException(e);
       }
 
-    } finally {
-      pluginManager.unlock();
     }
   }
 

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

@@ -0,0 +1,101 @@
+package software.amazon.jdbc.util;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.LockSupport;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class ResourceLockTest {
+
+    @Test
+    void testObtainClose() {
+        final ResourceLock lock = new ResourceLock();
+
+        assertFalse(lock.isLocked());
+        assertFalse(lock.isHeldByCurrentThread());
+
+        try (ResourceLock ignore = lock.obtain()) {
+            assertTrue(lock.isLocked());
+            assertTrue(lock.isHeldByCurrentThread());
+        }
+
+        assertFalse(lock.isLocked());
+        assertFalse(lock.isHeldByCurrentThread());
+    }
+
+    @Test
+    void testObtainWhenMultiThreadedExpectLinearExecution() throws InterruptedException, ExecutionException {
+        CallWithResourceLock callWithResourceLock = new CallWithResourceLock();
+
+        int levelOfConcurrency = 5;
+
+        ExecutorService executorService = Executors.newFixedThreadPool(levelOfConcurrency);
+        try {
+            List<Callable<CounterPair>> callables = new ArrayList<>();
+            for (int i = 0; i < levelOfConcurrency; i++) {
+                callables.add(callWithResourceLock::invoke);
+            }
+
+            // expect linear execution
+            List<Future<CounterPair>> results = executorService.invokeAll(callables);
+
+            Set<Integer> preLockSet = new HashSet<>();
+            Set<Integer> postLockSet = new HashSet<>();
+            for (Future<CounterPair> result : results) {
+                CounterPair entry = result.get();
+                preLockSet.add(entry.preLock);
+                postLockSet.add(entry.postLock);
+            }
+
+            assertEquals(levelOfConcurrency, postLockSet.size()); // linear execution inside resource lock block
+            assertEquals(1, preLockSet.size()); // all threads called invoke before any finish
+
+        } finally {
+            executorService.shutdown();
+        }
+    }
+
+    static final class CallWithResourceLock {
+
+        // wait enough time to allow concurrent threads to block on the lock
+        final long waitTime = TimeUnit.MILLISECONDS.toNanos(20);
+        final ResourceLock lock = new ResourceLock();
+        final AtomicInteger counter = new AtomicInteger();
+
+        /**
+         * Invoke returning 'pre lock' and 'post lock' counter value.
+         */
+        CounterPair invoke() {
+            int preLock = counter.get();
+            try (ResourceLock ignore = lock.obtain()) {
+                int postLock = counter.get();
+                LockSupport.parkNanos(waitTime);
+                counter.incrementAndGet();
+                return new CounterPair(preLock, postLock);
+            }
+        }
+    }
+
+    static final class CounterPair {
+        final int preLock;
+        final int postLock;
+
+        CounterPair(int preLock, int postLock) {
+            this.preLock = preLock;
+            this.postLock = postLock;
+        }
+    }
+
+}

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

@@ -49,11 +49,11 @@ void init() {
     doAnswer(invocation -> {
       pluginManagerLock.lock();
       return null;
-    }).when(pluginManager).lock();
+    }).when(pluginManager).acquireLock();
     doAnswer(invocation -> {
       pluginManagerLock.unlock();
       return null;
-    }).when(pluginManager).unlock();
+    }).when(pluginManager).releaseLock();
 
     doAnswer(invocation -> {
       boolean lockIsFree = testLock.tryLock();