From 7f43c62cb60855bd2184eb2472b4212fe4e16653 Mon Sep 17 00:00:00 2001
From: DigitalSonic <digitalsonic.dxf@gmail.com>
Date: Tue, 4 Feb 2020 17:01:39 +0800
Subject: [PATCH] add a NodeRegister to register Node to ZK

---
 pom.xml                                       |  13 +-
 .../druid/pool/ha/node/FileNodeListener.java  |   3 +
 .../druid/pool/ha/node/ZookeeperNodeInfo.java |  74 +++++++++
 .../pool/ha/node/ZookeeperNodeRegister.java   | 156 ++++++++++++++++++
 .../ha/node/ZookeeperNodeRegisterTest.java    | 103 ++++++++++++
 5 files changed, 348 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeInfo.java
 create mode 100644 src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegister.java
 create mode 100644 src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegisterTest.java

diff --git a/pom.xml b/pom.xml
index 4d4db226ae..def015ee6d 100644
--- a/pom.xml
+++ b/pom.xml
@@ -527,7 +527,12 @@
 			<version>2.5.3</version>
 			<scope>provided</scope>
 		</dependency>
-
+		<dependency>
+			<groupId>org.apache.curator</groupId>
+			<artifactId>curator-recipes</artifactId>
+			<version>4.2.0</version>
+			<scope>provided</scope>
+		</dependency>
 
 		<!--
 			test dependences
@@ -695,5 +700,11 @@
             <systemPath>${project.basedir}/lib/openjdk-1.8-jconsole.jar</systemPath>
         </dependency>
 
+		<dependency>
+			<groupId>org.apache.curator</groupId>
+			<artifactId>curator-test</artifactId>
+			<version>4.2.0</version>
+			<scope>test</scope>
+		</dependency>
     </dependencies>
 </project>
diff --git a/src/main/java/com/alibaba/druid/pool/ha/node/FileNodeListener.java b/src/main/java/com/alibaba/druid/pool/ha/node/FileNodeListener.java
index 2bc1b9cef2..17e20a1e0f 100644
--- a/src/main/java/com/alibaba/druid/pool/ha/node/FileNodeListener.java
+++ b/src/main/java/com/alibaba/druid/pool/ha/node/FileNodeListener.java
@@ -111,6 +111,9 @@ public class FileNodeListener extends NodeListener {
         return events;
     }
 
+    /**
+     * Close the ScheduledExecutorService.
+     */
     @Override
     public void destroy() {
         if (executor == null || executor.isShutdown()) {
diff --git a/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeInfo.java b/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeInfo.java
new file mode 100644
index 0000000000..d049d8f533
--- /dev/null
+++ b/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeInfo.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright 1999-2020 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.druid.pool.ha.node;
+
+/**
+ * The information that will be added to an ephemeral node.
+ *
+ * @author DigitalSonic
+ */
+public class ZookeeperNodeInfo {
+    private String prefix = "";
+    private String host;
+    private Integer port;
+    private String username;
+    private String password;
+
+    public void setPrefix(String prefix) {
+        if (prefix != null && !prefix.trim().isEmpty()) {
+            this.prefix = prefix;
+            if (!prefix.endsWith(".")) {
+                this.prefix = prefix + ".";
+            }
+        }
+    }
+
+    public String getPrefix() {
+        return prefix;
+    }
+
+    public String getHost() {
+        return host;
+    }
+
+    public void setHost(String host) {
+        this.host = host;
+    }
+
+    public Integer getPort() {
+        return port;
+    }
+
+    public void setPort(Integer port) {
+        this.port = port;
+    }
+
+    public String getUsername() {
+        return username;
+    }
+
+    public void setUsername(String username) {
+        this.username = username;
+    }
+
+    public String getPassword() {
+        return password;
+    }
+
+    public void setPassword(String password) {
+        this.password = password;
+    }
+}
diff --git a/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegister.java b/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegister.java
new file mode 100644
index 0000000000..e609adeaa5
--- /dev/null
+++ b/src/main/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegister.java
@@ -0,0 +1,156 @@
+/*
+ * Copyright 1999-2020 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.druid.pool.ha.node;
+
+import com.alibaba.druid.support.logging.Log;
+import com.alibaba.druid.support.logging.LogFactory;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.CuratorFrameworkFactory;
+import org.apache.curator.framework.recipes.nodes.GroupMember;
+import org.apache.curator.retry.ExponentialBackoffRetry;
+
+import java.io.IOException;
+import java.io.StringWriter;
+import java.util.List;
+import java.util.Properties;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * A register which is used to register a node to an ephemeral node.
+ *
+ * @author DigitalSonic
+ */
+public class ZookeeperNodeRegister {
+    private final static Log LOG = LogFactory.getLog(ZookeeperNodeRegister.class);
+    private String zkConnectString;
+    private String path = "/ha-druid-datasources";
+    private CuratorFramework client;
+    private GroupMember member;
+    private Lock lock = new ReentrantLock();
+
+    /**
+     * Init a CuratorFramework if there's no CuratorFramework provided.
+     */
+    public void init() {
+        if (client == null) {
+            client = CuratorFrameworkFactory.builder()
+                    .connectionTimeoutMs(3000)
+                    .connectString(zkConnectString)
+                    .retryPolicy(new ExponentialBackoffRetry(5000, 3, 30000))
+                    .sessionTimeoutMs(30000)
+                    .build();
+            client.start();
+        }
+    }
+
+    /**
+     * Register a Node which has a Properties as the payload.
+     * <pre>
+     * CAUTION: only one node can be registered,
+     *          if you want to register another one,
+     *          call deregister first
+     * </pre>
+     *
+     * @param payload The information used to generate the payload Properties
+     * @return true, register successfully; false, skip the registeration
+     */
+    public boolean register(String nodeId, List<ZookeeperNodeInfo> payload) {
+        if (payload == null || payload.isEmpty()) {
+            return false;
+        }
+        try {
+            lock.lock();
+            if (member != null) {
+                LOG.warn("GroupMember has already registered. Please deregister first.");
+                return false;
+            }
+            String payloadString = getPropertiesString(payload);
+            member = new GroupMember(client, path, nodeId, payloadString.getBytes());
+            member.start();
+            LOG.info("Register Node["+ nodeId + "] in path[" + path + "].");
+            return true;
+        } finally {
+            lock.unlock();
+        }
+    }
+
+    /**
+     * Close the current GroupMember.
+     */
+    public void deregister() {
+        if (member != null) {
+            member.close();
+            member = null;
+        }
+    }
+
+    /**
+     * @see #deregister()
+     */
+    public void destroy() {
+        deregister();
+    }
+
+    private String getPropertiesString(List<ZookeeperNodeInfo> payload) {
+        Properties properties = new Properties();
+        for (ZookeeperNodeInfo n : payload) {
+            if (n.getHost() != null) {
+                properties.setProperty(n.getPrefix() + "host", n.getHost());
+            }
+            if (n.getPort() != null) {
+                properties.setProperty(n.getPrefix() + "port", n.getPort().toString());
+            }
+            if (n.getUsername() != null) {
+                properties.setProperty(n.getPrefix() + "username", n.getUsername());
+            }
+            if (n.getPassword() != null) {
+                properties.setProperty(n.getPrefix() + "password", n.getPassword());
+            }
+        }
+        StringWriter sw = new StringWriter();
+        try {
+            properties.store(sw, "");
+        } catch (IOException e) {
+            LOG.error("Why Properties.store goes wrong?", e);
+        }
+        return sw.toString();
+    }
+
+    public String getZkConnectString() {
+        return zkConnectString;
+    }
+
+    public void setZkConnectString(String zkConnectString) {
+        this.zkConnectString = zkConnectString;
+    }
+
+    public String getPath() {
+        return path;
+    }
+
+    public void setPath(String path) {
+        this.path = path;
+    }
+
+    public CuratorFramework getClient() {
+        return client;
+    }
+
+    public void setClient(CuratorFramework client) {
+        this.client = client;
+    }
+}
diff --git a/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegisterTest.java b/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegisterTest.java
new file mode 100644
index 0000000000..7c927961ed
--- /dev/null
+++ b/src/test/java/com/alibaba/druid/pool/ha/node/ZookeeperNodeRegisterTest.java
@@ -0,0 +1,103 @@
+package com.alibaba.druid.pool.ha.node;
+
+import com.alibaba.druid.support.logging.Log;
+import com.alibaba.druid.support.logging.LogFactory;
+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.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.io.StringReader;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Properties;
+
+import static org.junit.Assert.*;
+
+public class ZookeeperNodeRegisterTest {
+    private final static Log LOG = LogFactory.getLog(ZookeeperNodeRegisterTest.class);
+    private final String PATH = "/ha-druid-datasource";
+    private static TestingServer server;
+    private ZookeeperNodeRegister register;
+
+    @BeforeClass
+    public static void init() throws Exception {
+        server = new TestingServer();
+    }
+
+    @AfterClass
+    public static void destroy() throws Exception {
+        server.close();
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        register = new ZookeeperNodeRegister();
+        register.setZkConnectString(server.getConnectString());
+        register.setPath(PATH);
+        register.init();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        register.destroy();
+        register.getClient().close();
+    }
+
+    @Test
+    public void testRegister() throws Exception {
+        List<ZookeeperNodeInfo> payload = new ArrayList<ZookeeperNodeInfo>();
+        ZookeeperNodeInfo node1 = new ZookeeperNodeInfo();
+        node1.setPrefix("foo");
+        node1.setHost("127.0.0.1");
+        node1.setPort(1234);
+        node1.setUsername("foo");
+        node1.setPassword("password");
+        payload.add(node1);
+
+        ZookeeperNodeInfo node2 = new ZookeeperNodeInfo();
+        node2.setPrefix("bar");
+        node2.setHost("127.0.0.1");
+        node2.setPort(5678);
+        node2.setUsername("bar");
+        node2.setPassword("password");
+        payload.add(node2);
+
+        assertFalse(register.register("test-foo", null));
+        assertTrue(register.register("test-foo", payload));
+        assertFalse(register.register("test-foo", payload));
+
+        Thread.sleep(1000); // Wait for the node to be created.
+
+        CuratorFramework client = register.getClient();
+        List<String> children = client.getChildren().forPath(PATH);
+        assertFalse(children.isEmpty());
+        assertEquals(1, children.size());
+        assertEquals("test-foo", children.get(0));
+
+        byte[] bytes = client.getData().forPath(PATH + "/test-foo");
+        Properties properties = new Properties();
+        String str = new String(bytes);
+        LOG.info("ZK Data: " + str);
+        properties.load(new StringReader(str));
+        validateNodeProperties(node1, properties);
+        validateNodeProperties(node2, properties);
+    }
+
+    private void validateNodeProperties(ZookeeperNodeInfo node, Properties properties) {
+        assertTrue(properties.containsKey(node.getPrefix() + "host"));
+        assertTrue(properties.containsKey(node.getPrefix() + "port"));
+        assertTrue(properties.containsKey(node.getPrefix() + "username"));
+        assertTrue(properties.containsKey(node.getPrefix() + "password"));
+
+        assertEquals(node.getHost(), properties.getProperty(node.getPrefix() + "host"));
+        assertEquals(node.getPort().toString(), properties.getProperty(node.getPrefix() + "port"));
+        assertEquals(node.getUsername(), properties.getProperty(node.getPrefix() + "username"));
+        assertEquals(node.getPassword(), properties.getProperty(node.getPrefix() + "password"));
+    }
+}
\ No newline at end of file