init

sentinel/sentinel-extension/sentinel-datasource-zookeeper/README.md

@@ -0,0 +1,27 @@
+# Sentinel DataSource ZooKeeper
+
+Sentinel DataSource ZooKeeper provides integration with ZooKeeper so that ZooKeeper
+can be the dynamic rule data source of Sentinel. The data source uses push model (listener).
+
+To use Sentinel DataSource ZooKeeper, you should add the following dependency:
+
+```xml
+<dependency>
+    <groupId>com.alibaba.csp</groupId>
+    <artifactId>sentinel-datasource-zookeeper</artifactId>
+    <version>x.y.z</version>
+</dependency>
+```
+
+Then you can create an `ZookeeperDataSource` and register to rule managers.
+For instance:
+
+```java
+// `path` is the data path in ZooKeeper
+ReadableDataSource<String, List<FlowRule>> flowRuleDataSource = new ZookeeperDataSource<>(remoteAddress, path, source -> JSON.parseObject(source, new TypeReference<List<FlowRule>>() {}));
+FlowRuleManager.register2Property(flowRuleDataSource.getProperty());
+```
+
+> Note: It's not recommended to add a large amount of rules to a single path (has limitation, also leads to bad performance).
+
+We've also provided an example: [sentinel-demo-zookeeper-datasource](https://github.com/alibaba/Sentinel/tree/master/sentinel-demo/sentinel-demo-zookeeper-datasource).

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

@@ -0,0 +1,72 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <artifactId>sentinel-extension</artifactId>
+        <groupId>com.alibaba.csp</groupId>
+        <version>1.8.3</version>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>sentinel-datasource-zookeeper</artifactId>
+    <packaging>jar</packaging>
+
+    <properties>
+        <zookeeper.version>3.4.14</zookeeper.version>
+        <curator.version>4.0.1</curator.version>
+        <curator.test.version>2.12.0</curator.test.version>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.alibaba.csp</groupId>
+            <artifactId>sentinel-datasource-extension</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>org.apache.zookeeper</groupId>
+            <artifactId>zookeeper</artifactId>
+            <version>${zookeeper.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.curator</groupId>
+            <artifactId>curator-recipes</artifactId>
+            <version>${curator.version}</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.apache.zookeeper</groupId>
+                    <artifactId>zookeeper</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.awaitility</groupId>
+            <artifactId>awaitility</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/sentinel-extension/sentinel-datasource-zookeeper/src/main/java/com/alibaba/csp/sentinel/datasource/zookeeper/ZookeeperDataSource.java

@@ -0,0 +1,216 @@
+package com.alibaba.csp.sentinel.datasource.zookeeper;
+
+import com.alibaba.csp.sentinel.concurrent.NamedThreadFactory;
+import com.alibaba.csp.sentinel.datasource.AbstractDataSource;
+import com.alibaba.csp.sentinel.datasource.Converter;
+import com.alibaba.csp.sentinel.log.RecordLog;
+import com.alibaba.csp.sentinel.util.StringUtil;
+import org.apache.curator.framework.AuthInfo;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.framework.recipes.cache.ChildData;
+import org.apache.curator.framework.recipes.cache.NodeCache;
+import org.apache.curator.framework.recipes.cache.NodeCacheListener;
+import org.apache.curator.retry.ExponentialBackoffRetry;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * A read-only {@code DataSource} with ZooKeeper backend.
+ *
+ * @author guonanjun
+ */
+public class ZookeeperDataSource<T> extends AbstractDataSource<String, T> {
+
+    private static final int RETRY_TIMES = 3;
+    private static final int SLEEP_TIME = 1000;
+
+    private static volatile Map<String, CuratorFramework> zkClientMap = new HashMap<>();
+    private static final Object lock = new Object();
+
+
+    private final ExecutorService pool = new ThreadPoolExecutor(1, 1, 0, TimeUnit.MILLISECONDS,
+            new ArrayBlockingQueue<Runnable>(1), new NamedThreadFactory("sentinel-zookeeper-ds-update", true),
+            new ThreadPoolExecutor.DiscardOldestPolicy());
+
+    private NodeCacheListener listener;
+    private final String path;
+
+    private CuratorFramework zkClient = null;
+    private NodeCache nodeCache = null;
+
+    public ZookeeperDataSource(final String serverAddr, final String path, Converter<String, T> parser) {
+        super(parser);
+        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(path)) {
+            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], path=[%s]", serverAddr, path));
+        }
+        this.path = path;
+
+        init(serverAddr, null);
+    }
+
+    /**
+     * This constructor is Nacos-style.
+     */
+    public ZookeeperDataSource(final String serverAddr, final String groupId, final String dataId,
+                               Converter<String, T> parser) {
+        super(parser);
+        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(groupId) || StringUtil.isBlank(dataId)) {
+            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], groupId=[%s], dataId=[%s]", serverAddr, groupId, dataId));
+        }
+        this.path = getPath(groupId, dataId);
+
+        init(serverAddr, null);
+    }
+
+    /**
+     * This constructor adds authentication information.
+     */
+    public ZookeeperDataSource(final String serverAddr, final List<AuthInfo> authInfos, final String groupId, final String dataId,
+                               Converter<String, T> parser) {
+        super(parser);
+        if (StringUtil.isBlank(serverAddr) || StringUtil.isBlank(groupId) || StringUtil.isBlank(dataId)) {
+            throw new IllegalArgumentException(String.format("Bad argument: serverAddr=[%s], authInfos=[%s], groupId=[%s], dataId=[%s]", serverAddr, authInfos, groupId, dataId));
+        }
+        this.path = getPath(groupId, dataId);
+
+        init(serverAddr, authInfos);
+    }
+
+    private void init(final String serverAddr, final List<AuthInfo> authInfos) {
+        initZookeeperListener(serverAddr, authInfos);
+        loadInitialConfig();
+    }
+
+    private void loadInitialConfig() {
+        try {
+            T newValue = loadConfig();
+            if (newValue == null) {
+                RecordLog.warn("[ZookeeperDataSource] WARN: initial config is null, you may have to check your data source");
+            }
+            getProperty().updateValue(newValue);
+        } catch (Exception ex) {
+            RecordLog.warn("[ZookeeperDataSource] Error when loading initial config", ex);
+        }
+    }
+
+    private void initZookeeperListener(final String serverAddr, final List<AuthInfo> authInfos) {
+        try {
+
+            this.listener = new NodeCacheListener() {
+                @Override
+                public void nodeChanged() {
+
+                    try {
+                        T newValue = loadConfig();
+                        RecordLog.info("[ZookeeperDataSource] New property value received for ({}, {}): {}",
+                            serverAddr, path, newValue);
+                        // Update the new value to the property.
+                        getProperty().updateValue(newValue);
+                    } catch (Exception ex) {
+                        RecordLog.warn("[ZookeeperDataSource] loadConfig exception", ex);
+                    }
+                }
+            };
+
+            String zkKey = getZkKey(serverAddr, authInfos);
+            if (zkClientMap.containsKey(zkKey)) {
+                this.zkClient = zkClientMap.get(zkKey);
+            } else {
+                synchronized (lock) {
+                    if (!zkClientMap.containsKey(zkKey)) {
+                        CuratorFramework zc = null;
+                        if (authInfos == null || authInfos.size() == 0) {
+                            zc = CuratorFrameworkFactory.newClient(serverAddr, new ExponentialBackoffRetry(SLEEP_TIME, RETRY_TIMES));
+                        } else {
+                            zc = CuratorFrameworkFactory.builder().
+                                    connectString(serverAddr).
+                                    retryPolicy(new ExponentialBackoffRetry(SLEEP_TIME, RETRY_TIMES)).
+                                    authorization(authInfos).
+                                    build();
+                        }
+                        this.zkClient = zc;
+                        this.zkClient.start();
+                        Map<String, CuratorFramework> newZkClientMap = new HashMap<>(zkClientMap.size());
+                        newZkClientMap.putAll(zkClientMap);
+                        newZkClientMap.put(zkKey, zc);
+                        zkClientMap = newZkClientMap;
+                    } else {
+                        this.zkClient = zkClientMap.get(zkKey);
+                    }
+                }
+            }
+
+            this.nodeCache = new NodeCache(this.zkClient, this.path);
+            this.nodeCache.getListenable().addListener(this.listener, this.pool);
+            this.nodeCache.start();
+        } catch (Exception e) {
+            RecordLog.warn("[ZookeeperDataSource] Error occurred when initializing Zookeeper data source", e);
+            e.printStackTrace();
+        }
+    }
+
+    @Override
+    public String readSource() throws Exception {
+        if (this.zkClient == null) {
+            throw new IllegalStateException("Zookeeper has not been initialized or error occurred");
+        }
+        String configInfo = null;
+        ChildData childData = nodeCache.getCurrentData();
+        if (null != childData && childData.getData() != null) {
+
+            configInfo = new String(childData.getData());
+        }
+        return configInfo;
+    }
+
+    @Override
+    public void close() throws Exception {
+        if (this.nodeCache != null) {
+            this.nodeCache.getListenable().removeListener(listener);
+            this.nodeCache.close();
+        }
+        if (this.zkClient != null) {
+            this.zkClient.close();
+        }
+        pool.shutdown();
+    }
+
+    private String getPath(String groupId, String dataId) {
+        return String.format("/%s/%s", groupId, dataId);
+    }
+
+    private String getZkKey(final String serverAddr, final List<AuthInfo> authInfos) {
+        if (authInfos == null || authInfos.size() == 0) {
+            return serverAddr;
+        }
+        StringBuilder builder = new StringBuilder(64);
+        builder.append(serverAddr).append(getAuthInfosKey(authInfos));
+        return builder.toString();
+    }
+
+    private String getAuthInfosKey(List<AuthInfo> authInfos) {
+        StringBuilder builder = new StringBuilder(32);
+        for (AuthInfo authInfo : authInfos) {
+            if (authInfo == null) {
+                builder.append("{}");
+            } else {
+                builder.append("{" + "sc=" + authInfo.getScheme() + ",au=" + Arrays.toString(authInfo.getAuth()) + "}");
+            }
+        }
+        return builder.toString();
+    }
+
+    protected CuratorFramework getZkClient() {
+        return this.zkClient;
+    }
+
+
+}

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

@@ -0,0 +1,224 @@
+package com.alibaba.csp.sentinel.datasource.zookeeper;
+
+import com.alibaba.csp.sentinel.datasource.Converter;
+import com.alibaba.csp.sentinel.datasource.ReadableDataSource;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.degrade.DegradeRule;
+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.AuthInfo;
+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.ZooDefs;
+import org.apache.zookeeper.data.ACL;
+import org.apache.zookeeper.data.Id;
+import org.apache.zookeeper.data.Stat;
+import org.apache.zookeeper.server.auth.DigestAuthenticationProvider;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @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";
+
+        ReadableDataSource<String, List<FlowRule>> flowRuleDataSource = new ZookeeperDataSource<List<FlowRule>>(remoteAddress, path,
+                new Converter<String, List<FlowRule>>() {
+                    @Override
+                    public List<FlowRule> convert(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();
+    }
+
+
+
+
+    @Test
+    public void testZooKeeperDataSourceAuthorization() throws Exception {
+        TestingServer server = new TestingServer(21812);
+        server.start();
+
+        final String remoteAddress = server.getConnectString();
+        final String groupId = "sentinel-zk-ds-demo";
+        final String dataId = "flow-HK";
+        final String path = "/" + groupId + "/" + dataId;
+        final String scheme = "digest";
+        final String auth = "root:123456";
+
+        AuthInfo authInfo = new AuthInfo(scheme, auth.getBytes());
+        List<AuthInfo> authInfoList = Collections.singletonList(authInfo);
+
+        CuratorFramework zkClient = CuratorFrameworkFactory.builder().
+                connectString(remoteAddress).
+                retryPolicy(new ExponentialBackoffRetry(3, 100)).
+                authorization(authInfoList).
+                build();
+        zkClient.start();
+        Stat stat = zkClient.checkExists().forPath(path);
+        if (stat == null) {
+            ACL acl = new ACL(ZooDefs.Perms.ALL, new Id(scheme, DigestAuthenticationProvider.generateDigest(auth)));
+            zkClient.create().creatingParentContainersIfNeeded().withACL(Collections.singletonList(acl)).forPath(path, null);
+        }
+
+        ReadableDataSource<String, List<FlowRule>> flowRuleDataSource = new ZookeeperDataSource<List<FlowRule>>(remoteAddress,
+                authInfoList, groupId, dataId,
+                new Converter<String, List<FlowRule>>() {
+                    @Override
+                    public List<FlowRule> convert(String source) {
+                        return JSON.parseObject(source, new TypeReference<List<FlowRule>>() {
+                        });
+                    }
+                });
+        FlowRuleManager.register2Property(flowRuleDataSource.getProperty());
+
+
+        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());
+
+        await().timeout(5, TimeUnit.SECONDS)
+                .until(new Callable<Boolean>() {
+                    @Override
+                    public Boolean call() throws Exception {
+                        List<FlowRule> rules = FlowRuleManager.getRules();
+                        return rules != null && !rules.isEmpty();
+                    }
+                });
+
+        List<FlowRule> rules = FlowRuleManager.getRules();
+        boolean exists = false;
+        for (FlowRule r : rules) {
+            if (resourceName.equals(r.getResource())) {
+                exists = true;
+                assertEquals(count, new Double(r.getCount()).longValue());
+            }
+        }
+        assertTrue(exists);
+    }
+
+
+    /**
+     * Test whether different dataSources can share the same zkClient when the connection parameters are the same.
+     * @throws Exception
+     */
+    @Test
+    public void testZooKeeperDataSourceSameZkClient() throws Exception {
+        TestingServer server = new TestingServer(21813);
+        server.start();
+
+        final String remoteAddress = server.getConnectString();
+        final String flowPath = "/sentinel-zk-ds-demo/flow-HK";
+        final String degradePath = "/sentinel-zk-ds-demo/degrade-HK";
+
+
+        ZookeeperDataSource<List<FlowRule>> flowRuleZkDataSource = new ZookeeperDataSource<>(remoteAddress, flowPath,
+                new Converter<String, List<FlowRule>>() {
+                    @Override
+                    public List<FlowRule> convert(String source) {
+                        return JSON.parseObject(source, new TypeReference<List<FlowRule>>() {
+                        });
+                    }
+                });
+        ZookeeperDataSource<List<DegradeRule>> degradeRuleZkDataSource = new ZookeeperDataSource<>(remoteAddress, degradePath,
+                new Converter<String, List<DegradeRule>>() {
+                    @Override
+                    public List<DegradeRule> convert(String source) {
+                        return JSON.parseObject(source, new TypeReference<List<DegradeRule>>() {
+                        });
+                    }
+                });
+
+
+        Assert.assertTrue(flowRuleZkDataSource.getZkClient() == degradeRuleZkDataSource.getZkClient());
+
+
+        final String groupId = "sentinel-zk-ds-demo";
+        final String flowDataId = "flow-HK";
+        final String degradeDataId = "degrade-HK";
+        final String scheme = "digest";
+        final String auth = "root:123456";
+        AuthInfo authInfo = new AuthInfo(scheme, auth.getBytes());
+        List<AuthInfo> authInfoList = Collections.singletonList(authInfo);
+
+
+        ZookeeperDataSource<List<FlowRule>> flowRuleZkAutoDataSource = new ZookeeperDataSource<List<FlowRule>>(remoteAddress,
+                authInfoList, groupId, flowDataId,
+                new Converter<String, List<FlowRule>>() {
+                    @Override
+                    public List<FlowRule> convert(String source) {
+                        return JSON.parseObject(source, new TypeReference<List<FlowRule>>() {
+                        });
+                    }
+                });
+
+        ZookeeperDataSource<List<DegradeRule>> degradeRuleZkAutoDataSource = new ZookeeperDataSource<List<DegradeRule>>(remoteAddress,
+                authInfoList, groupId, degradeDataId,
+                new Converter<String, List<DegradeRule>>() {
+                    @Override
+                    public List<DegradeRule> convert(String source) {
+                        return JSON.parseObject(source, new TypeReference<List<DegradeRule>>() {
+                        });
+                    }
+                });
+
+        Assert.assertTrue(flowRuleZkAutoDataSource.getZkClient() == degradeRuleZkAutoDataSource.getZkClient());
+
+        server.stop();
+    }
+
+
+
+}