Add custom connection validation to ConnectionPoolSupport #3081 (#3138)

* Add custom connection validation to ConnectionPoolSupport

* Polishing

---------

Co-authored-by: Tihomir Mateev <[email protected]>

Author: big-cir

Diff for: 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 @@
  * </pre>
  *
  * @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,30 +81,67 @@ private ConnectionPoolSupport() {
      */
     public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
             Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> 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 <T> connection type.
+     * @return the connection pool.
+     */
+    public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
+            Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config, Predicate<T> 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 <T> connection type.
      * @return the connection pool.
      */
     @SuppressWarnings("unchecked")
     public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
             Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> 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 <T> connection type.
+     * @return the connection pool.
+     */
+    @SuppressWarnings("unchecked")
+    public static <T extends StatefulConnection<?, ?>> GenericObjectPool<T> createGenericObjectPool(
+            Supplier<T> connectionSupplier, GenericObjectPoolConfig<T> config, boolean wrapConnections,
+            Predicate<T> 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<Origin<T>> poolRef = new AtomicReference<>();
 
-        GenericObjectPool<T> pool = new GenericObjectPool<T>(new RedisPooledObjectFactory<T>(connectionSupplier), config) {
+        GenericObjectPool<T> pool = new GenericObjectPool<T>(
+                new RedisPooledObjectFactory<>(connectionSupplier, validationPredicate), config) {
 
             @Override
             public T borrowObject() throws Exception {
@@ -144,20 +184,38 @@ 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 <T> connection type.
      * @return the connection pool.
      */
     @SuppressWarnings("unchecked")
     public static <T extends StatefulConnection<?, ?>> SoftReferenceObjectPool<T> createSoftReferenceObjectPool(
             Supplier<T> 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 <T> connection type.
+     * @return the connection pool.
+     */
+    @SuppressWarnings("unchecked")
+    public static <T extends StatefulConnection<?, ?>> SoftReferenceObjectPool<T> createSoftReferenceObjectPool(
+            Supplier<T> connectionSupplier, boolean wrapConnections, Predicate<T> validationPredicate) {
 
         LettuceAssert.notNull(connectionSupplier, "Connection supplier must not be null");
 
         AtomicReference<Origin<T>> poolRef = new AtomicReference<>();
 
-        SoftReferenceObjectPool<T> pool = new SoftReferenceObjectPool<T>(new RedisPooledObjectFactory<>(connectionSupplier)) {
+        SoftReferenceObjectPool<T> pool = new SoftReferenceObjectPool<T>(
+                new RedisPooledObjectFactory<>(connectionSupplier, validationPredicate)) {
 
             private final Lock lock = new ReentrantLock();
 
@@ -200,8 +258,11 @@ private static class RedisPooledObjectFactory<T extends StatefulConnection<?, ?>
 
         private final Supplier<T> connectionSupplier;
 
-        RedisPooledObjectFactory(Supplier<T> connectionSupplier) {
+        private final Predicate<T> validationPredicate;
+
+        RedisPooledObjectFactory(Supplier<T> connectionSupplier, Predicate<T> validationPredicate) {
             this.connectionSupplier = connectionSupplier;
+            this.validationPredicate = validationPredicate;
         }
 
         @Override
@@ -221,7 +282,7 @@ public PooledObject<T> wrap(T obj) {
 
         @Override
         public boolean validateObject(PooledObject<T> p) {
-            return p.getObject().isOpen();
+            return this.validationPredicate.test(p.getObject());
         }
 
     }

Diff for: 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<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
+                .createGenericObjectPool(client::connect, new GenericObjectPoolConfig<>(), connection -> {
+                    try {
+                        return "PONG".equals(connection.sync().ping());
+                    } catch (Exception e) {
+                        return false;
+                    }
+                });
+
+        try (StatefulRedisConnection<String, String> connection = pool.borrowObject()) {
+
+            RedisCommandFactory factory = new RedisCommandFactory(connection);
+            SimpleCommands commands = factory.getCommands(SimpleCommands.class);
+            commands.get("foo");
+        }
+
+        pool.close();
+    }
+
     @Test
     void shouldWorkWithAsyncPooledConnection() {
 

Diff for: 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<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
+                .createGenericObjectPool(() -> client.connect(), new GenericObjectPoolConfig<>(), false, connection -> {
+                    try {
+                        return "PONG".equals(connection.sync().ping());
+                    } catch (Exception e) {
+                        return false;
+                    }
+                });
+
+        borrowAndReturn(pool);
+
+        StatefulRedisConnection<String, String> 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<StatefulRedisConnection<String, String>> pool = ConnectionPoolSupport
+                .createSoftReferenceObjectPool(() -> client.connect(), false, connection -> {
+                    try {
+                        return "PONG".equals(connection.sync().ping());
+                    } catch (Exception e) {
+                        return false;
+                    }
+                });
+
+        borrowAndReturn(pool);
+
+        StatefulRedisConnection<String, String> connection = pool.borrowObject();
+        assertThat(Proxy.isProxyClass(connection.getClass())).isFalse();
+        pool.returnObject(connection);
+
+        connection.close();
+        pool.close();
+    }
+
     @Test
     void genericPoolUsingWrappingShouldPropagateExceptionsCorrectly() throws Exception {