Add support for ReadFrom.subnet redis#1569

Lettuce now supports ReadFrom.subnet which is an option of reading from any node in the given subnets.

src/main/java/io/lettuce/core/ReadFrom.java

@@ -26,6 +26,7 @@
  * @author Mark Paluch
  * @author Ryosuke Hasebe
  * @author Omer Cilingir
+ * @author Yohei Ueki
  * @since 4.0
  */
 public abstract class ReadFrom {
@@ -110,6 +111,17 @@ public abstract class ReadFrom {
      */
     public static final ReadFrom ANY_REPLICA = new ReadFromImpl.ReadFromAnyReplica();
 
+    /**
+     * Setting to read from any node in the subnets.
+     *
+     * @param cidrNotations CIDR-block notation strings, e.g., "192.168.0.0/16".
+     * @return an instance of {@link ReadFromImpl.ReadFromSubnet}.
+     * @since x.x.x
+     */
+    public static ReadFrom subnet(String... cidrNotations) {
+        return new ReadFromImpl.ReadFromSubnet(cidrNotations);
+    }
+
     /**
      * Chooses the nodes from the matching Redis nodes that match this read selector.
      *
@@ -178,6 +190,10 @@ public static ReadFrom valueOf(String name) {
             return ANY_REPLICA;
         }
 
+        if (name.equalsIgnoreCase("subnet")) {
+            throw new IllegalArgumentException("subnet must be created via ReadFrom#subnet");
+        }
+
         throw new IllegalArgumentException("ReadFrom " + name + " not supported");
     }
 

src/main/java/io/lettuce/core/ReadFromImpl.java

@@ -15,19 +15,26 @@
  */
 package io.lettuce.core;
 
+import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Predicate;
 
+import io.lettuce.core.internal.LettuceAssert;
 import io.lettuce.core.internal.LettuceLists;
+import io.lettuce.core.internal.LettuceStrings;
 import io.lettuce.core.models.role.RedisNodeDescription;
+import io.netty.handler.ipfilter.IpFilterRuleType;
+import io.netty.handler.ipfilter.IpSubnetFilterRule;
+import io.netty.util.internal.SocketUtils;
 
 /**
  * Collection of common read setting implementations.
  *
  * @author Mark Paluch
  * @author Omer Cilingir
+ * @author Yohei Ueki
  * @since 4.0
  */
 class ReadFromImpl {
@@ -129,6 +136,60 @@ public ReadFromAnyReplica() {
 
     }
 
+    /**
+     * Read from any node in the subnets.
+     *
+     * @since x.x.x
+     */
+    static final class ReadFromSubnet extends ReadFrom {
+
+        private final List<IpSubnetFilterRule> rules = new ArrayList<>();
+
+        /**
+         * @param cidrNotations CIDR-block notation strings, e.g., "192.168.0.0/16".
+         */
+        ReadFromSubnet(String... cidrNotations) {
+            LettuceAssert.notEmpty(cidrNotations, "cidrNotations must not be empty");
+
+            for (String cidrNotation : cidrNotations) {
+                // parts[0]: ipAddress (e.g., "192.168.0.0")
+                // parts[1]: cidrPrefix (e.g., "16")
+                String[] parts = cidrNotation.split("/");
+                LettuceAssert.isTrue(parts.length == 2, "cidrNotation must have exact one '/'");
+
+                rules.add(new IpSubnetFilterRule(parts[0], Integer.parseInt(parts[1]), IpFilterRuleType.ACCEPT));
+            }
+        }
+
+        @Override
+        public List<RedisNodeDescription> select(Nodes nodes) {
+            List<RedisNodeDescription> result = new ArrayList<>(nodes.getNodes().size());
+            for (RedisNodeDescription node : nodes) {
+                if (isInSubnet(node.getUri())) {
+                    result.add(node);
+                }
+            }
+
+            return result;
+        }
+
+        private boolean isInSubnet(RedisURI redisURI) {
+            if (LettuceStrings.isEmpty(redisURI.getHost())) {
+                return false;
+            }
+
+            InetSocketAddress address = SocketUtils.socketAddress(redisURI.getHost(), redisURI.getPort());
+
+            for (IpSubnetFilterRule rule : rules) {
+                if (rule.matches(address)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+    }
+
     /**
      * {@link Predicate}-based {@link ReadFrom} implementation.
      *

src/test/java/io/lettuce/core/cluster/ReadFromUnitTests.java

@@ -18,6 +18,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -26,6 +27,7 @@
 import org.junit.jupiter.api.Test;
 
 import io.lettuce.core.ReadFrom;
+import io.lettuce.core.RedisURI;
 import io.lettuce.core.cluster.models.partitions.Partitions;
 import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
 import io.lettuce.core.models.role.RedisNodeDescription;
@@ -34,6 +36,7 @@
  * @author Mark Paluch
  * @author Ryosuke Hasebe
  * @author Omer Cilingir
+ * @author Yohei Ueki
  */
 class ReadFromUnitTests {
 
@@ -90,6 +93,50 @@ void anyReplica() {
         assertThat(result).hasSize(2).containsExactly(nearest, replica);
     }
 
+    @Test
+    void subnet() {
+        ReadFrom sut = ReadFrom.subnet("192.0.2.0/24", "2001:db8:abcd:0000::/52");
+
+        RedisClusterNode nodeInSubnetIpv4 = createNodeWithHost("192.0.2.1");
+        RedisClusterNode nodeNotInSubnetIpv4 = createNodeWithHost("198.51.100.1");
+        RedisClusterNode nodeInSubnetIpv6 = createNodeWithHost("2001:db8:abcd:0000::1");
+        RedisClusterNode nodeNotInSubnetIpv6 = createNodeWithHost("2001:db8:abcd:1000::");
+        RedisClusterNode nodeHostname = createNodeWithHost("example.com");
+
+        List<RedisNodeDescription> result = sut
+                .select(getNodes(nodeInSubnetIpv4, nodeNotInSubnetIpv4, nodeInSubnetIpv6, nodeNotInSubnetIpv6, nodeHostname));
+
+        assertThat(result).hasSize(2).containsExactly(nodeInSubnetIpv4, nodeInSubnetIpv6);
+    }
+
+    @Test
+    void subnetLocalhost() {
+        ReadFrom sut = ReadFrom.subnet("127.0.0.0/8", "::1/128");
+
+        RedisClusterNode localhost = createNodeWithHost("localhost");
+
+        List<RedisNodeDescription> result = sut.select(getNodes(localhost));
+
+        assertThat(result).hasSize(1).containsExactly(localhost);
+    }
+
+    private RedisClusterNode createNodeWithHost(String host) {
+        RedisClusterNode node = new RedisClusterNode();
+        node.setUri(RedisURI.Builder.redis(host).build());
+        return node;
+    }
+
+    @Test
+    void subnetInvalidCidr() {
+        // malformed CIDR notation
+        assertThatThrownBy(() -> ReadFrom.subnet("192.0.2.1//1")).isInstanceOf(IllegalArgumentException.class);
+        // malformed ipAddress
+        assertThatThrownBy(() -> ReadFrom.subnet("foo.bar/12")).isInstanceOf(IllegalArgumentException.class);
+        // malformed cidrPrefix
+        assertThatThrownBy(() -> ReadFrom.subnet("192.0.2.1/40")).isInstanceOf(IllegalArgumentException.class);
+        assertThatThrownBy(() -> ReadFrom.subnet("192.0.2.1/foo")).isInstanceOf(IllegalArgumentException.class);
+    }
+
     @Test
     void valueOfNull() {
         assertThatThrownBy(() -> ReadFrom.valueOf(null)).isInstanceOf(IllegalArgumentException.class);
@@ -130,6 +177,11 @@ void valueOfAnyReplica() {
         assertThat(ReadFrom.valueOf("anyReplica")).isEqualTo(ReadFrom.ANY_REPLICA);
     }
 
+    @Test
+    void valueOfSubnet() {
+        assertThatThrownBy(() -> ReadFrom.valueOf("subnet")).isInstanceOf(IllegalArgumentException.class);
+    }
+
     private ReadFrom.Nodes getNodes() {
         return new ReadFrom.Nodes() {
             @Override
@@ -144,4 +196,22 @@ public Iterator<RedisNodeDescription> iterator() {
         };
 
     }
+
+    private ReadFrom.Nodes getNodes(RedisNodeDescription... nodes) {
+        return new ReadFrom.Nodes() {
+
+            @Override
+            public List<RedisNodeDescription> getNodes() {
+                return Arrays.asList(nodes);
+            }
+
+            @Override
+            public Iterator<RedisNodeDescription> iterator() {
+                return getNodes().iterator();
+            }
+
+        };
+
+    }
+
 }

src/test/java/io/lettuce/core/cluster/RedisClusterReadFromIntegrationTests.java

@@ -32,6 +32,7 @@
 
 /**
  * @author Mark Paluch
+ * @author Yohei Ueki
  */
 @SuppressWarnings("unchecked")
 @ExtendWith(LettuceExtension.class)
@@ -112,4 +113,16 @@ void readWriteNearest() {
         connection.getConnection(ClusterTestSettings.host, ClusterTestSettings.port2).sync().waitForReplication(1, 1000);
         assertThat(sync.get(key)).isEqualTo("value1");
     }
+
+    @Test
+    void readWriteSubnet() {
+
+        connection.setReadFrom(ReadFrom.subnet("0.0.0.0/0", "::/0"));
+
+        sync.set(key, "value1");
+
+        connection.getConnection(ClusterTestSettings.host, ClusterTestSettings.port2).sync().waitForReplication(1, 1000);
+        assertThat(sync.get(key)).isEqualTo("value1");
+    }
+
 }