diff --git a/src/main/java/io/lettuce/core/support/ConnectionPoolSupport.java b/src/main/java/io/lettuce/core/support/ConnectionPoolSupport.java index f3ab8cb85..a048dfe27 100644 --- a/src/main/java/io/lettuce/core/support/ConnectionPoolSupport.java +++ b/src/main/java/io/lettuce/core/support/ConnectionPoolSupport.java @@ -6,6 +6,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Predicate; import java.util.function.Supplier; import org.apache.commons.pool2.BasePooledObjectFactory; @@ -60,6 +61,7 @@ * * * @author Mark Paluch + * @author dae won * @since 4.3 */ public abstract class ConnectionPoolSupport { @@ -69,7 +71,8 @@ private ConnectionPoolSupport() { /** * Creates a new {@link GenericObjectPool} using the {@link Supplier}. Allocated instances are wrapped and must not be - * returned with {@link ObjectPool#returnObject(Object)}. + * returned with {@link ObjectPool#returnObject(Object)}. By default, connections are validated by checking their + * {@link StatefulConnection#isOpen()} method. * * @param connectionSupplier must not be {@code null}. * @param config must not be {@code null}. @@ -78,16 +81,32 @@ private ConnectionPoolSupport() { */ public static > GenericObjectPool createGenericObjectPool( Supplier connectionSupplier, GenericObjectPoolConfig config) { - return createGenericObjectPool(connectionSupplier, config, true); + return createGenericObjectPool(connectionSupplier, config, true, (c) -> c.isOpen()); } /** - * Creates a new {@link GenericObjectPool} using the {@link Supplier}. + * Creates a new {@link GenericObjectPool} using the {@link Supplier}. Allocated instances are wrapped and must not be + * returned with {@link ObjectPool#returnObject(Object)}. + * + * @param connectionSupplier must not be {@code null}. + * @param config must not be {@code null}. + * @param validationPredicate a {@link Predicate} to help validate connections + * @param connection type. + * @return the connection pool. + */ + public static > GenericObjectPool createGenericObjectPool( + Supplier connectionSupplier, GenericObjectPoolConfig config, Predicate validationPredicate) { + return createGenericObjectPool(connectionSupplier, config, true, validationPredicate); + } + + /** + * Creates a new {@link GenericObjectPool} using the {@link Supplier}. By default, connections are validated by checking + * their {@link StatefulConnection#isOpen()} method. * * @param connectionSupplier must not be {@code null}. * @param config must not be {@code null}. * @param wrapConnections {@code false} to return direct connections that need to be returned to the pool using - * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connection that are returned to the pool + * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connections that are returned to the pool * when invoking {@link StatefulConnection#close()}. * @param connection type. * @return the connection pool. @@ -95,13 +114,34 @@ private ConnectionPoolSupport() { @SuppressWarnings("unchecked") public static > GenericObjectPool createGenericObjectPool( Supplier connectionSupplier, GenericObjectPoolConfig config, boolean wrapConnections) { + return createGenericObjectPool(connectionSupplier, config, wrapConnections, (c) -> c.isOpen()); + } + + /** + * Creates a new {@link GenericObjectPool} using the {@link Supplier}. + * + * @param connectionSupplier must not be {@code null}. + * @param config must not be {@code null}. + * @param wrapConnections {@code false} to return direct connections that need to be returned to the pool using + * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connections that are returned to the pool + * when invoking {@link StatefulConnection#close()}. + * @param validationPredicate a {@link Predicate} to help validate connections + * @param connection type. + * @return the connection pool. + */ + @SuppressWarnings("unchecked") + public static > GenericObjectPool createGenericObjectPool( + Supplier connectionSupplier, GenericObjectPoolConfig config, boolean wrapConnections, + Predicate validationPredicate) { LettuceAssert.notNull(connectionSupplier, "Connection supplier must not be null"); LettuceAssert.notNull(config, "GenericObjectPoolConfig must not be null"); + LettuceAssert.notNull(validationPredicate, "Connection validator must not be null"); AtomicReference> poolRef = new AtomicReference<>(); - GenericObjectPool pool = new GenericObjectPool(new RedisPooledObjectFactory(connectionSupplier), config) { + GenericObjectPool pool = new GenericObjectPool( + new RedisPooledObjectFactory<>(connectionSupplier, validationPredicate), config) { @Override public T borrowObject() throws Exception { @@ -144,7 +184,7 @@ public void returnObject(T obj) { * * @param connectionSupplier must not be {@code null}. * @param wrapConnections {@code false} to return direct connections that need to be returned to the pool using - * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connection that are returned to the pool + * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connections that are returned to the pool * when invoking {@link StatefulConnection#close()}. * @param connection type. * @return the connection pool. @@ -152,12 +192,30 @@ public void returnObject(T obj) { @SuppressWarnings("unchecked") public static > SoftReferenceObjectPool createSoftReferenceObjectPool( Supplier connectionSupplier, boolean wrapConnections) { + return createSoftReferenceObjectPool(connectionSupplier, wrapConnections, (c) -> c.isOpen()); + } + + /** + * Creates a new {@link SoftReferenceObjectPool} using the {@link Supplier}. + * + * @param connectionSupplier must not be {@code null}. + * @param wrapConnections {@code false} to return direct connections that need to be returned to the pool using + * {@link ObjectPool#returnObject(Object)}. {@code true} to return wrapped connections that are returned to the pool + * when invoking {@link StatefulConnection#close()}. + * @param validationPredicate a {@link Predicate} to help validate connections + * @param connection type. + * @return the connection pool. + */ + @SuppressWarnings("unchecked") + public static > SoftReferenceObjectPool createSoftReferenceObjectPool( + Supplier connectionSupplier, boolean wrapConnections, Predicate validationPredicate) { LettuceAssert.notNull(connectionSupplier, "Connection supplier must not be null"); AtomicReference> poolRef = new AtomicReference<>(); - SoftReferenceObjectPool pool = new SoftReferenceObjectPool(new RedisPooledObjectFactory<>(connectionSupplier)) { + SoftReferenceObjectPool pool = new SoftReferenceObjectPool( + new RedisPooledObjectFactory<>(connectionSupplier, validationPredicate)) { private final Lock lock = new ReentrantLock(); @@ -200,8 +258,11 @@ private static class RedisPooledObjectFactory private final Supplier connectionSupplier; - RedisPooledObjectFactory(Supplier connectionSupplier) { + private final Predicate validationPredicate; + + RedisPooledObjectFactory(Supplier connectionSupplier, Predicate validationPredicate) { this.connectionSupplier = connectionSupplier; + this.validationPredicate = validationPredicate; } @Override @@ -221,7 +282,7 @@ public PooledObject wrap(T obj) { @Override public boolean validateObject(PooledObject p) { - return p.getObject().isOpen(); + return this.validationPredicate.test(p.getObject()); } } diff --git a/src/test/java/io/lettuce/core/dynamic/RedisCommandsIntegrationTests.java b/src/test/java/io/lettuce/core/dynamic/RedisCommandsIntegrationTests.java index 21c2b1e98..aa0098796 100644 --- a/src/test/java/io/lettuce/core/dynamic/RedisCommandsIntegrationTests.java +++ b/src/test/java/io/lettuce/core/dynamic/RedisCommandsIntegrationTests.java @@ -112,6 +112,28 @@ void shouldWorkWithPooledConnection() throws Exception { pool.close(); } + @Test + void shouldWorkWithPooledConnectionAndCustomValidation() throws Exception { + + GenericObjectPool> pool = ConnectionPoolSupport + .createGenericObjectPool(client::connect, new GenericObjectPoolConfig<>(), connection -> { + try { + return "PONG".equals(connection.sync().ping()); + } catch (Exception e) { + return false; + } + }); + + try (StatefulRedisConnection connection = pool.borrowObject()) { + + RedisCommandFactory factory = new RedisCommandFactory(connection); + SimpleCommands commands = factory.getCommands(SimpleCommands.class); + commands.get("foo"); + } + + pool.close(); + } + @Test void shouldWorkWithAsyncPooledConnection() { diff --git a/src/test/java/io/lettuce/core/support/ConnectionPoolSupportIntegrationTests.java b/src/test/java/io/lettuce/core/support/ConnectionPoolSupportIntegrationTests.java index ed38435ed..db64d7c33 100644 --- a/src/test/java/io/lettuce/core/support/ConnectionPoolSupportIntegrationTests.java +++ b/src/test/java/io/lettuce/core/support/ConnectionPoolSupportIntegrationTests.java @@ -131,6 +131,27 @@ void genericPoolShouldWorkWithPlainConnections() throws Exception { pool.close(); } + @Test + void genericPoolShouldWorkWithValidationPredicate() throws Exception { + + GenericObjectPool> pool = ConnectionPoolSupport + .createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false, connection -> { + try { + return "PONG".equals(connection.sync().ping()); + } catch (Exception e) { + return false; + } + }); + + borrowAndReturn(pool); + + StatefulRedisConnection connection = pool.borrowObject(); + assertThat(Proxy.isProxyClass(connection.getClass())).isFalse(); + pool.returnObject(connection); + + pool.close(); + } + @Test void softReferencePoolShouldWorkWithPlainConnections() throws Exception { @@ -147,6 +168,28 @@ void softReferencePoolShouldWorkWithPlainConnections() throws Exception { pool.close(); } + @Test + void softReferencePoolShouldWorkWithValidationPredicate() throws Exception { + + SoftReferenceObjectPool> pool = ConnectionPoolSupport + .createSoftReferenceObjectPool(() -> client.connect(), false, connection -> { + try { + return "PONG".equals(connection.sync().ping()); + } catch (Exception e) { + return false; + } + }); + + borrowAndReturn(pool); + + StatefulRedisConnection connection = pool.borrowObject(); + assertThat(Proxy.isProxyClass(connection.getClass())).isFalse(); + pool.returnObject(connection); + + connection.close(); + pool.close(); + } + @Test void genericPoolUsingWrappingShouldPropagateExceptionsCorrectly() throws Exception {