[helix] branch master updated: Fix ZkAddress null case for TaskStateModelFactory (#1627)

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hulee pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/helix.git


The following commit(s) were added to refs/heads/master by this push:
     new 2ce66cb  Fix ZkAddress null case for TaskStateModelFactory (#1627)
2ce66cb is described below

commit 2ce66cb7f5a33c9a0161f11b30833057604823ba
Author: Neal Sun <ne...@gmail.com>
AuthorDate: Tue Feb 2 17:57:23 2021 -0800

    Fix ZkAddress null case for TaskStateModelFactory (#1627)
    
    If zkAddress is null for the HelixManager passed to TaskStateModelFactory and multizk is not enabled, a NullPointerException will be raised. This PR fixes that by adding a null check.
---
 .../org/apache/helix/task/TaskStateModelFactory.java     |  5 +++--
 .../org/apache/helix/task/TestTaskStateModelFactory.java | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java b/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java
index 215c895..292c26a 100644
--- a/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java
+++ b/helix-core/src/main/java/org/apache/helix/task/TaskStateModelFactory.java
@@ -142,8 +142,9 @@ public class TaskStateModelFactory extends StateModelFactory<TaskStateModel> {
     // want to enforce it, which may cause backward incompatibility.
     RealmAwareZkClient.RealmAwareZkClientConfig clientConfig =
         new RealmAwareZkClient.RealmAwareZkClientConfig().setZkSerializer(new ZNRecordSerializer());
+    String zkAddress = manager.getMetadataStoreConnectionString();
 
-    if (Boolean.getBoolean(SystemPropertyKeys.MULTI_ZK_ENABLED)) {
+    if (Boolean.getBoolean(SystemPropertyKeys.MULTI_ZK_ENABLED) || zkAddress == null) {
       String clusterName = manager.getClusterName();
       String shardingKey = HelixUtil.clusterNameToShardingKey(clusterName);
       RealmAwareZkClient.RealmAwareZkConnectionConfig connectionConfig =
@@ -158,7 +159,7 @@ public class TaskStateModelFactory extends StateModelFactory<TaskStateModel> {
     }
 
     return SharedZkClientFactory.getInstance().buildZkClient(
-        new HelixZkClient.ZkConnectionConfig(manager.getMetadataStoreConnectionString()),
+        new HelixZkClient.ZkConnectionConfig(zkAddress),
         clientConfig.createHelixZkClientConfig().setZkSerializer(new ZNRecordSerializer()));
   }
 }
diff --git a/helix-core/src/test/java/org/apache/helix/task/TestTaskStateModelFactory.java b/helix-core/src/test/java/org/apache/helix/task/TestTaskStateModelFactory.java
index 3919cde..dc56792 100644
--- a/helix-core/src/test/java/org/apache/helix/task/TestTaskStateModelFactory.java
+++ b/helix-core/src/test/java/org/apache/helix/task/TestTaskStateModelFactory.java
@@ -24,6 +24,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
+import org.apache.helix.HelixManager;
 import org.apache.helix.SystemPropertyKeys;
 import org.apache.helix.integration.manager.MockParticipantManager;
 import org.apache.helix.integration.task.TaskTestBase;
@@ -31,10 +32,14 @@ import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.msdcommon.constant.MetadataStoreRoutingConstants;
 import org.apache.helix.msdcommon.mock.MockMetadataStoreDirectoryServer;
 import org.apache.helix.zookeeper.api.client.RealmAwareZkClient;
+import org.apache.helix.zookeeper.impl.client.FederatedZkClient;
 import org.apache.helix.zookeeper.routing.RoutingDataManager;
+import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import static org.mockito.Mockito.when;
+
 
 public class TestTaskStateModelFactory extends TaskTestBase {
   // This value has to be different from the default value to verify correctness
@@ -85,6 +90,17 @@ public class TestTaskStateModelFactory extends TaskTestBase {
     Assert.assertEquals(TaskUtil
         .getTargetThreadPoolSize(zkClient, anyParticipantManager.getClusterName(),
             anyParticipantManager.getInstanceName()), TEST_TARGET_TASK_THREAD_POOL_SIZE);
+    Assert.assertTrue(zkClient instanceof FederatedZkClient);
+
+    // Turn off multiZk mode in System config, and remove zkAddress
+    System.setProperty(SystemPropertyKeys.MULTI_ZK_ENABLED, "false");
+    HelixManager participantManager = Mockito.spy(anyParticipantManager);
+    when(participantManager.getMetadataStoreConnectionString()).thenReturn(null);
+    zkClient = TaskStateModelFactory.createZkClient(participantManager);
+    Assert.assertEquals(TaskUtil
+        .getTargetThreadPoolSize(zkClient, anyParticipantManager.getClusterName(),
+            anyParticipantManager.getInstanceName()), TEST_TARGET_TASK_THREAD_POOL_SIZE);
+    Assert.assertTrue(zkClient instanceof FederatedZkClient);
 
     // Restore system properties
     if (prevMultiZkEnabled == null) {