From e4df01a0330451c964e7326dd85bfe2c11821cd1 Mon Sep 17 00:00:00 2001
From: Wenjun Ruan <wenjun@apache.org>
Date: Tue, 21 May 2024 10:32:15 +0800
Subject: [PATCH] Increase block-until-connected in ZookeeperRegistryTestCase
 (#16041)

---
 .../plugin/registry/zookeeper/ZookeeperRegistry.java | 12 +++++++-----
 .../src/test/resources/application.yaml              |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/main/java/org/apache/dolphinscheduler/plugin/registry/zookeeper/ZookeeperRegistry.java b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/main/java/org/apache/dolphinscheduler/plugin/registry/zookeeper/ZookeeperRegistry.java
index 67f505d79b..76313df02a 100644
--- a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/main/java/org/apache/dolphinscheduler/plugin/registry/zookeeper/ZookeeperRegistry.java
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/main/java/org/apache/dolphinscheduler/plugin/registry/zookeeper/ZookeeperRegistry.java
@@ -25,6 +25,7 @@ import org.apache.dolphinscheduler.registry.api.Registry;
 import org.apache.dolphinscheduler.registry.api.RegistryException;
 import org.apache.dolphinscheduler.registry.api.SubscribeListener;
 
+import org.apache.commons.lang3.time.DurationUtils;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.framework.api.ACLProvider;
@@ -76,8 +77,8 @@ final class ZookeeperRegistry implements Registry {
                         .connectString(properties.getConnectString())
                         .retryPolicy(retryPolicy)
                         .namespace(properties.getNamespace())
-                        .sessionTimeoutMs((int) properties.getSessionTimeout().toMillis())
-                        .connectionTimeoutMs((int) properties.getConnectionTimeout().toMillis());
+                        .sessionTimeoutMs(DurationUtils.toMillisInt(properties.getSessionTimeout()))
+                        .connectionTimeoutMs(DurationUtils.toMillisInt(properties.getConnectionTimeout()));
 
         final String digest = properties.getDigest();
         if (!Strings.isNullOrEmpty(digest)) {
@@ -102,9 +103,10 @@ final class ZookeeperRegistry implements Registry {
     public void start() {
         client.start();
         try {
-            if (!client.blockUntilConnected((int) properties.getBlockUntilConnected().toMillis(), MILLISECONDS)) {
+            if (!client.blockUntilConnected(DurationUtils.toMillisInt(properties.getBlockUntilConnected()),
+                    MILLISECONDS)) {
                 client.close();
-                throw new RegistryException("zookeeper connect timeout: " + properties.getConnectString());
+                throw new RegistryException("zookeeper connect failed in : " + properties.getConnectString() + "ms");
             }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
@@ -120,7 +122,7 @@ final class ZookeeperRegistry implements Registry {
     @Override
     public void connectUntilTimeout(@NonNull Duration timeout) throws RegistryException {
         try {
-            if (!client.blockUntilConnected((int) timeout.toMillis(), MILLISECONDS)) {
+            if (!client.blockUntilConnected(DurationUtils.toMillisInt(timeout), MILLISECONDS)) {
                 throw new RegistryException(
                         String.format("Cannot connect to registry in %s s", timeout.getSeconds()));
             }
diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/test/resources/application.yaml b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/test/resources/application.yaml
index 7c272e55d5..92902a608c 100644
--- a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/test/resources/application.yaml
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-zookeeper/src/test/resources/application.yaml
@@ -26,5 +26,5 @@ registry:
       max-retries: 5
     session-timeout: 30s
     connection-timeout: 9s
-    block-until-connected: 600ms
+    block-until-connected: 3s
     digest: ~