From d87d2d80dd88db76aa2d99f4f34ecfeeda584afe Mon Sep 17 00:00:00 2001
From: tswstarplanet <tswstarplanet@apache.org>
Date: Thu, 9 Jul 2020 20:27:41 +0800
Subject: [PATCH] add maxWaitTime config when build zookeeper client (#3133)

* add maxWaitTime config when build zookeeper client

* throw exception when connect zk timeout

* use dedicated exception instead of a generic one

Co-authored-by: dailidong <dailidong66@gmail.com>
Co-authored-by: qiaozhanwei <qiaozhanwei@outlook.com>
---
 .../dolphinscheduler/service/zk/ZookeeperConfig.java  | 11 +++++++++++
 .../service/zk/ZookeeperOperator.java                 |  5 ++++-
 .../src/main/resources/zookeeper.properties           |  3 ++-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperConfig.java b/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperConfig.java
index 5bdc6f8cd..57ac13e3b 100644
--- a/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperConfig.java
+++ b/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperConfig.java
@@ -52,6 +52,9 @@ public class ZookeeperConfig {
     @Value("${zookeeper.dolphinscheduler.root:/dolphinscheduler}")
     private String dsRoot;
 
+    @Value("${zookeeper.max.wait.time:10000}")
+    private int maxWaitTime;
+
     public String getServerList() {
         return serverList;
     }
@@ -115,4 +118,12 @@ public class ZookeeperConfig {
     public void setDsRoot(String dsRoot) {
         this.dsRoot = dsRoot;
     }
+
+    public int getMaxWaitTime() {
+        return maxWaitTime;
+    }
+
+    public void setMaxWaitTime(int maxWaitTime) {
+        this.maxWaitTime = maxWaitTime;
+    }
 }
\ No newline at end of file
diff --git a/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperOperator.java b/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperOperator.java
index cd479d44c..ba3a3bfec 100644
--- a/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperOperator.java
+++ b/dolphinscheduler-service/src/main/java/org/apache/dolphinscheduler/service/zk/ZookeeperOperator.java
@@ -37,6 +37,7 @@ import org.springframework.stereotype.Component;
 
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.dolphinscheduler.common.utils.Preconditions.checkNotNull;
 
@@ -109,7 +110,9 @@ public class ZookeeperOperator implements InitializingBean {
         zkClient = builder.build();
         zkClient.start();
         try {
-            zkClient.blockUntilConnected();
+            if (!zkClient.blockUntilConnected(zookeeperConfig.getMaxWaitTime(), TimeUnit.MILLISECONDS)) {
+                throw new IllegalStateException("Connect zookeeper expire max wait time");
+            }
         } catch (final Exception ex) {
             throw new RuntimeException(ex);
         }
diff --git a/dolphinscheduler-service/src/main/resources/zookeeper.properties b/dolphinscheduler-service/src/main/resources/zookeeper.properties
index 2204467ac..31f5f43c4 100644
--- a/dolphinscheduler-service/src/main/resources/zookeeper.properties
+++ b/dolphinscheduler-service/src/main/resources/zookeeper.properties
@@ -26,4 +26,5 @@ zookeeper.quorum=localhost:2181
 #zookeeper.connection.timeout=30000
 #zookeeper.retry.base.sleep=100
 #zookeeper.retry.max.sleep=30000
-#zookeeper.retry.maxtime=10
\ No newline at end of file
+#zookeeper.retry.maxtime=10
+#zookeeper.max.wait.time=10000
\ No newline at end of file
-- 
GitLab