Update ZooKeeper data source integration

- Code refinement
- Add test cases

Signed-off-by: Eric Zhao <sczyh16@gmail.com>

sentinel-demo/sentinel-demo-zookeeper-datasource/pom.xml

@@ -14,7 +14,6 @@
     <properties>
         <zookeeper.version>3.4.13</zookeeper.version>
         <curator.version>4.0.1</curator.version>
-        <curator-test.version>2.12.0</curator-test.version>
     </properties>
 
     <dependencies>
@@ -44,8 +43,8 @@
 
         <dependency>
             <groupId>org.apache.curator</groupId>
-            <artifactId>curator-test</artifactId>
+            <artifactId>curator-recipes</artifactId>
-            <version>${curator-test.version}</version>
+            <version>${curator.version}</version>
             <exclusions>
                 <exclusion>
                     <groupId>org.apache.zookeeper</groupId>

sentinel-demo/sentinel-demo-zookeeper-datasource/src/main/java/com/alibaba/csp/sentinel/demo/datasource/zookeeper/ZookeeperConfigSender.java

@@ -3,7 +3,6 @@ package com.alibaba.csp.sentinel.demo.datasource.zookeeper;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.ExponentialBackoffRetry;
-import org.apache.curator.test.TestingServer;
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.data.Stat;
 
@@ -18,11 +17,7 @@ public class ZookeeperConfigSender {
     private static final int SLEEP_TIME = 1000;
 
     public static void main(String[] args) throws Exception {
-
+        final String remoteAddress = "localhost:2181";
-        // 启动Zookeeper服务
-        TestingServer server = new TestingServer(2181);
-
-        final String remoteAddress = server.getConnectString();
         final String groupId = "Sentinel-Demo";
         final String dataId = "SYSTEM-CODE-DEMO-FLOW";
         final String rule = "[\n"
@@ -44,18 +39,14 @@ public class ZookeeperConfigSender {
             zkClient.create().creatingParentContainersIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path, null);
         }
         zkClient.setData().forPath(path, rule.getBytes());
-        // zkClient.delete().forPath(path);
 
         try {
-            Thread.sleep(30000L);
+            Thread.sleep(3000);
         } catch (InterruptedException e) {
             e.printStackTrace();
         }
 
         zkClient.close();
-
-        //停止zookeeper服务
-        server.stop();
     }
 
     private static String getPath(String groupId, String dataId) {

sentinel-extension/sentinel-datasource-zookeeper/pom.xml

@@ -15,6 +15,7 @@
     <properties>
         <zookeeper.version>3.4.13</zookeeper.version>
         <curator.version>4.0.1</curator.version>
+        <curator.test.version>2.12.0</curator.test.version>
     </properties>
 
     <dependencies>
@@ -39,5 +40,28 @@
                 </exclusion>
             </exclusions>
         </dependency>
+
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.curator</groupId>
+            <artifactId>curator-test</artifactId>
+            <version>${curator.test.version}</version>
+            <scope>test</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.apache.zookeeper</groupId>
+                    <artifactId>zookeeper</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+        <dependency>
+            <groupId>com.alibaba</groupId>
+            <artifactId>fastjson</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

sentinel-extension/sentinel-datasource-zookeeper/src/main/java/com/alibaba/csp/sentinel/datasource/zookeeper/ZookeeperDataSource.java

@@ -10,6 +10,7 @@ import com.alibaba.csp.sentinel.datasource.AbstractDataSource;
 import com.alibaba.csp.sentinel.datasource.ConfigParser;
 import com.alibaba.csp.sentinel.log.RecordLog;
 import com.alibaba.csp.sentinel.util.StringUtil;
+
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.framework.recipes.cache.ChildData;
@@ -34,36 +35,31 @@ public class ZookeeperDataSource<T> extends AbstractDataSource<String, T> {
         new ThreadPoolExecutor.DiscardOldestPolicy());
 
     private NodeCacheListener listener;
-    private final String groupId;
-    private final String dataId;
     private final String path;
 
     private CuratorFramework zkClient = null;
     private NodeCache nodeCache = null;
 
-    public ZookeeperDataSource(final String serverAddr, final String groupId, final String dataId,
+    public ZookeeperDataSource(final String serverAddr, final String path, ConfigParser<String, T> parser) {
-                               ConfigParser<String, T> parser) {
         super(parser);
-        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(groupId) || StringUtil.isBlank(dataId)) {
+        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(path)) {
-            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], groupId=[%s], dataId=[%s]",
+            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], path=[%s]", serverAddr, path));
-                    serverAddr, groupId, dataId));
         }
-        this.groupId = groupId;
+        this.path = path;
-        this.dataId = dataId;
-        this.path = getPath(groupId, dataId);
 
         init(serverAddr);
     }
 
-    public ZookeeperDataSource(final String serverAddr, final String path, ConfigParser<String, T> parser) {
+    /**
+     * This constructor is Nacos-style.
+     */
+    public ZookeeperDataSource(final String serverAddr, final String groupId, final String dataId,
+                               ConfigParser<String, T> parser) {
         super(parser);
-        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(path)) {
+        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(groupId) || StringUtil.isBlank(dataId)) {
-            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], path=[%s]",
+            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], groupId=[%s], dataId=[%s]", serverAddr, groupId, dataId));
-                    serverAddr, path));
         }
-        this.path = path;
+        this.path = getPath(groupId, dataId);
-        this.groupId = null;
-        this.dataId = null;
 
         init(serverAddr);
     }
@@ -90,7 +86,7 @@ public class ZookeeperDataSource<T> extends AbstractDataSource<String, T> {
 
             this.listener = new NodeCacheListener() {
                 @Override
-                public void nodeChanged() throws Exception {
+                public void nodeChanged() {
                     String configInfo = null;
                     ChildData childData = nodeCache.getCurrentData();
                     if (null != childData && childData.getData() != null) {
@@ -146,17 +142,6 @@ public class ZookeeperDataSource<T> extends AbstractDataSource<String, T> {
     }
 
     private String getPath(String groupId, String dataId) {
-        String path = "";
+        return String.format("/%s/%s", groupId, dataId);
-        if (groupId.startsWith("/")) {
-            path += groupId;
-        } else {
-            path += "/" + groupId;
-        }
-        if (dataId.startsWith("/")) {
-            path += dataId;
-        } else {
-            path += "/" + dataId;
-        }
-        return path;
     }
 }

sentinel-extension/sentinel-datasource-zookeeper/src/test/java/com/alibaba/csp/sentinel/datasource/zookeeper/ZookeeperDataSourceTest.java

@@ -0,0 +1,84 @@
+package com.alibaba.csp.sentinel.datasource.zookeeper;
+
+import java.util.Collections;
+import java.util.List;
+
+import com.alibaba.csp.sentinel.datasource.ConfigParser;
+import com.alibaba.csp.sentinel.datasource.DataSource;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
+import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.TypeReference;
+
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.retry.ExponentialBackoffRetry;
+import org.apache.curator.test.TestingServer;
+import org.apache.zookeeper.CreateMode;
+import org.apache.zookeeper.data.Stat;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+/**
+ * @author Eric Zhao
+ */
+public class ZookeeperDataSourceTest {
+
+    @Test
+    public void testZooKeeperDataSource() throws Exception {
+        TestingServer server = new TestingServer(21812);
+        server.start();
+
+        final String remoteAddress = server.getConnectString();
+        final String path = "/sentinel-zk-ds-demo/flow-HK";
+
+        DataSource<String, List<FlowRule>> flowRuleDataSource = new ZookeeperDataSource<List<FlowRule>>(remoteAddress, path,
+            new ConfigParser<String, List<FlowRule>>() {
+                @Override
+                public List<FlowRule> parse(String source) {
+                    return JSON.parseObject(source, new TypeReference<List<FlowRule>>() {});
+                }
+            });
+        FlowRuleManager.register2Property(flowRuleDataSource.getProperty());
+
+        CuratorFramework zkClient = CuratorFrameworkFactory.newClient(remoteAddress,
+            new ExponentialBackoffRetry(3, 1000));
+        zkClient.start();
+        Stat stat = zkClient.checkExists().forPath(path);
+        if (stat == null) {
+            zkClient.create().creatingParentContainersIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path, null);
+        }
+
+        final String resourceName = "HK";
+        publishThenTestFor(zkClient, path, resourceName, 10);
+        publishThenTestFor(zkClient, path, resourceName, 15);
+
+        zkClient.close();
+        server.stop();
+    }
+
+    private void publishThenTestFor(CuratorFramework zkClient, String path, String resourceName, long count) throws Exception {
+        FlowRule rule = new FlowRule().setResource(resourceName)
+            .setLimitApp("default")
+            .as(FlowRule.class)
+            .setCount(count)
+            .setGrade(RuleConstant.FLOW_GRADE_QPS);
+        String ruleString = JSON.toJSONString(Collections.singletonList(rule));
+        zkClient.setData().forPath(path, ruleString.getBytes());
+
+        Thread.sleep(5000);
+
+        List<FlowRule> rules = FlowRuleManager.getRules();
+        assertTrue(rules != null && !rules.isEmpty());
+        boolean exists = false;
+        for (FlowRule r : rules) {
+            if (resourceName.equals(r.getResource())) {
+                exists = true;
+                assertEquals(count, new Double(r.getCount()).longValue());
+            }
+        }
+        assertTrue(exists);
+    }
+}