From 9764fd7ba3fd971933e0e05d6e2fb78b7dc5aa5b Mon Sep 17 00:00:00 2001 From: Dave Cramer Date: Thu, 17 Nov 2022 09:39:11 -0500 Subject: [PATCH] Use try with resources to simplify locking code --- .../amazon/jdbc/ConnectionPluginManager.java | 15 +-- .../amazon/jdbc/util/ResourceLock.java | 54 ++++++++++ .../amazon/jdbc/util/WrapperUtils.java | 12 +-- .../amazon/jdbc/util/ResourceLockTest.java | 101 ++++++++++++++++++ .../amazon/jdbc/util/WrapperUtilsTest.java | 4 +- 5 files changed, 168 insertions(+), 18 deletions(-) create mode 100644 wrapper/src/main/java/software/amazon/jdbc/util/ResourceLock.java create mode 100644 wrapper/src/test/java/software/amazon/jdbc/util/ResourceLockTest.java diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java index 277cfcb12..0bddba8dc 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java +++ b/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 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(); } /** diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ResourceLock.java b/wrapper/src/main/java/software/amazon/jdbc/util/ResourceLock.java new file mode 100644 index 000000000..d5f1f7689 --- /dev/null +++ b/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. + * + *

Example use

+ * 
{@code
+ *
+ *   try (ResourceLock ignore = lock.obtain()) {
+ *     // do something while holding the resource lock
+ *   }
+ *
+ * }
+ */ +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(); + } +} \ No newline at end of file diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/WrapperUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/WrapperUtils.java index 207ec871e..6f2d55e8c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/WrapperUtils.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/WrapperUtils.java @@ -180,9 +180,8 @@ public static T executeWithPlugins( final JdbcCallable 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 executeWithPlugins( } catch (final InstantiationException e) { throw new RuntimeException(e); } - - } finally { - pluginManager.unlock(); } } @@ -216,9 +212,7 @@ public static 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 executeWithPlugins( throw new RuntimeException(e); } - } finally { - pluginManager.unlock(); } } diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/ResourceLockTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/ResourceLockTest.java new file mode 100644 index 000000000..3b43d6bb2 --- /dev/null +++ b/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> callables = new ArrayList<>(); + for (int i = 0; i < levelOfConcurrency; i++) { + callables.add(callWithResourceLock::invoke); + } + + // expect linear execution + List> results = executorService.invokeAll(callables); + + Set preLockSet = new HashSet<>(); + Set postLockSet = new HashSet<>(); + for (Future 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; + } + } + +} \ No newline at end of file diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/WrapperUtilsTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/WrapperUtilsTest.java index 814bdaae3..e7949e12a 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/WrapperUtilsTest.java +++ b/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();