Custom connection pool to MultiClusterPooledConnectionProvider (#3801)

* Add support for configuring connection pool to MultiClusterPooledConnectionProvider * Add ClusterConfig constructor to include ConnectionPoolConfig, add test * Update src/main/java/redis/clients/jedis/MultiClusterClientConfig.java Co-authored-by: M Sazzadul Hoque <[email protected]> * Update src/main/java/redis/clients/jedis/MultiClusterClientConfig.java Co-authored-by: M Sazzadul Hoque <[email protected]> * Update test * Update src/main/java/redis/clients/jedis/MultiClusterClientConfig.java Co-authored-by: M Sazzadul Hoque <[email protected]> * Update src/test/java/redis/clients/jedis/providers/MultiClusterPooledConnectionProviderTest.java Co-authored-by: M Sazzadul Hoque <[email protected]> * Format testConnectionPoolConfigApplied --------- Co-authored-by: M Sazzadul Hoque <[email protected]>
redis · May 20, 2024 · fd2c5e7 · fd2c5e7
1 parent 019776c
commit fd2c5e7
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 4 deletions.
diff --git a/src/main/java/redis/clients/jedis/MultiClusterClientConfig.java b/src/main/java/redis/clients/jedis/MultiClusterClientConfig.java
@@ -7,6 +7,7 @@
 import java.util.Arrays;
 import java.util.List;
 
+import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 import redis.clients.jedis.exceptions.JedisConnectionException;
 import redis.clients.jedis.exceptions.JedisValidationException;
 
@@ -175,12 +176,20 @@ public static class ClusterConfig {
         private int priority;
         private HostAndPort hostAndPort;
         private JedisClientConfig clientConfig;
+        private GenericObjectPoolConfig<Connection> connectionPoolConfig;
 
         public ClusterConfig(HostAndPort hostAndPort, JedisClientConfig clientConfig) {
             this.hostAndPort = hostAndPort;
             this.clientConfig = clientConfig;
         }
 
+        public ClusterConfig(HostAndPort hostAndPort, JedisClientConfig clientConfig,
+                             GenericObjectPoolConfig<Connection> connectionPoolConfig) {
+            this.hostAndPort = hostAndPort;
+            this.clientConfig = clientConfig;
+            this.connectionPoolConfig = connectionPoolConfig;
+        }
+
         public int getPriority() {
             return priority;
         }
@@ -196,6 +205,10 @@ public HostAndPort getHostAndPort() {
         public JedisClientConfig getJedisClientConfig() {
             return clientConfig;
         }
+
+        public GenericObjectPoolConfig<Connection> getConnectionPoolConfig() {
+            return connectionPoolConfig;
+        }
     }
 
     public static class Builder {

diff --git a/src/main/java/redis/clients/jedis/providers/MultiClusterPooledConnectionProvider.java b/src/main/java/redis/clients/jedis/providers/MultiClusterPooledConnectionProvider.java
@@ -14,6 +14,8 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
+
+import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -110,6 +112,7 @@ public MultiClusterPooledConnectionProvider(MultiClusterClientConfig multiCluste
 
         ClusterConfig[] clusterConfigs = multiClusterClientConfig.getClusterConfigs();
         for (ClusterConfig config : clusterConfigs) {
+            GenericObjectPoolConfig<Connection> poolConfig = config.getConnectionPoolConfig();
 
             String clusterId = "cluster:" + config.getPriority() + ":" + config.getHostAndPort();
 
@@ -128,9 +131,15 @@ public MultiClusterPooledConnectionProvider(MultiClusterClientConfig multiCluste
             circuitBreakerEventPublisher.onSlowCallRateExceeded(event -> log.error(String.valueOf(event)));
             circuitBreakerEventPublisher.onStateTransition(event -> log.warn(String.valueOf(event)));
 
-            multiClusterMap.put(config.getPriority(),
-                    new Cluster(new ConnectionPool(config.getHostAndPort(),
-                            config.getJedisClientConfig()), retry, circuitBreaker));
+            if (poolConfig != null) {
+                multiClusterMap.put(config.getPriority(),
+                        new Cluster(new ConnectionPool(config.getHostAndPort(),
+                                config.getJedisClientConfig(), poolConfig), retry, circuitBreaker));
+            } else {
+                multiClusterMap.put(config.getPriority(),
+                        new Cluster(new ConnectionPool(config.getHostAndPort(),
+                                config.getJedisClientConfig()), retry, circuitBreaker));
+            }
         }
 
         /// --- ///

diff --git a/src/test/java/redis/clients/jedis/providers/MultiClusterPooledConnectionProviderTest.java b/src/test/java/redis/clients/jedis/providers/MultiClusterPooledConnectionProviderTest.java
@@ -131,4 +131,22 @@ public void testSetActiveMultiClusterIndexOutOfRange() {
         provider.setActiveMultiClusterIndex(3); // Should throw an exception
     }
 
-}
+    @Test
+    public void testConnectionPoolConfigApplied() {
+        ConnectionPoolConfig poolConfig = new ConnectionPoolConfig();
+        poolConfig.setMaxTotal(8);
+        poolConfig.setMaxIdle(4);
+        poolConfig.setMinIdle(1);
+        ClusterConfig[] clusterConfigs = new ClusterConfig[2];
+        clusterConfigs[0] = new ClusterConfig(hostAndPort1, DefaultJedisClientConfig.builder().build(), poolConfig);
+        clusterConfigs[1] = new ClusterConfig(hostAndPort2, DefaultJedisClientConfig.builder().build(), poolConfig);
+        try (MultiClusterPooledConnectionProvider customProvider = new MultiClusterPooledConnectionProvider(
+                new MultiClusterClientConfig.Builder(clusterConfigs).build())) {
+            MultiClusterPooledConnectionProvider.Cluster activeCluster = customProvider.getCluster();
+            ConnectionPool connectionPool = activeCluster.getConnectionPool();
+            assertEquals(8, connectionPool.getMaxTotal());
+            assertEquals(4, connectionPool.getMaxIdle());
+            assertEquals(1, connectionPool.getMinIdle());
+        }
+    }
+}