[helix] branch master updated: Make HelixCustomCodeRunner multi-ZK aware (#1243)

[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 a86b7d4  Make HelixCustomCodeRunner multi-ZK aware (#1243)
a86b7d4 is described below

commit a86b7d46bea2743592baa2f2fede2c0f45bfacd6
Author: Hunter Lee <hu...@linkedin.com>
AuthorDate: Mon Aug 10 17:40:44 2020 -0700

    Make HelixCustomCodeRunner multi-ZK aware (#1243)
    
    This class was an edge case that was neglected as part of ZooScalability API migration. This commit makes HelixCustomCodeRunner realm-aware by adding a constructor that accepts a RealmAwareZkConnectionConfig.
---
 .../helix/participant/HelixCustomCodeRunner.java   | 38 ++++++++++++++++++----
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/participant/HelixCustomCodeRunner.java b/helix-core/src/main/java/org/apache/helix/participant/HelixCustomCodeRunner.java
index 711a584..9afdef6 100644
--- a/helix-core/src/main/java/org/apache/helix/participant/HelixCustomCodeRunner.java
+++ b/helix-core/src/main/java/org/apache/helix/participant/HelixCustomCodeRunner.java
@@ -23,16 +23,19 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
+import org.apache.helix.SystemPropertyKeys;
 import org.apache.helix.zookeeper.api.client.HelixZkClient;
 import org.apache.helix.HelixConstants.ChangeType;
 import org.apache.helix.HelixDataAccessor;
 import org.apache.helix.HelixManager;
 import org.apache.helix.PropertyKey.Builder;
 import org.apache.helix.manager.zk.ZKHelixDataAccessor;
-import org.apache.helix.manager.zk.ZNRecordSerializer;
 import org.apache.helix.manager.zk.ZkBaseDataAccessor;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.model.IdealState.RebalanceMode;
+import org.apache.helix.zookeeper.api.client.RealmAwareZkClient;
+import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
+import org.apache.helix.zookeeper.impl.client.FederatedZkClient;
 import org.apache.helix.zookeeper.impl.factory.SharedZkClientFactory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -63,6 +66,7 @@ public class HelixCustomCodeRunner {
   private String _resourceName;
   private final HelixManager _manager;
   private final String _zkAddr;
+  private final RealmAwareZkClient.RealmAwareZkConnectionConfig _connectionConfig;
   private GenericLeaderStandbyStateModelFactory _stateModelFty;
 
   /**
@@ -73,6 +77,19 @@ public class HelixCustomCodeRunner {
   public HelixCustomCodeRunner(HelixManager manager, String zkAddr) {
     _manager = manager;
     _zkAddr = zkAddr;
+    _connectionConfig = null;
+  }
+
+  /**
+   * Constructs a HelixCustomCodeRunner that will be run on multi-zk mode.
+   * @param manager
+   * @param connectionConfig config with a multi-realm fields set
+   */
+  public HelixCustomCodeRunner(HelixManager manager,
+      RealmAwareZkClient.RealmAwareZkConnectionConfig connectionConfig) {
+    _manager = manager;
+    _zkAddr = null;
+    _connectionConfig = connectionConfig;
   }
 
   /**
@@ -130,14 +147,23 @@ public class HelixCustomCodeRunner {
 
     StateMachineEngine stateMach = _manager.getStateMachineEngine();
     stateMach.registerStateModelFactory(LEADER_STANDBY, _stateModelFty, _resourceName);
-    HelixZkClient zkClient = null;
+    RealmAwareZkClient zkClient = null;
     try {
       // manually add ideal state for participant leader using LeaderStandby
       // model
-      HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();
-      clientConfig.setZkSerializer(new ZNRecordSerializer());
-      zkClient = SharedZkClientFactory.getInstance()
-          .buildZkClient(new HelixZkClient.ZkConnectionConfig(_zkAddr), clientConfig);
+      if (Boolean.getBoolean(SystemPropertyKeys.MULTI_ZK_ENABLED) || _zkAddr == null) {
+        // Use multi-zk mode (FederatedZkClient)
+        RealmAwareZkClient.RealmAwareZkClientConfig clientConfig =
+            new RealmAwareZkClient.RealmAwareZkClientConfig();
+        clientConfig.setZkSerializer(new ZNRecordSerializer());
+        zkClient = new FederatedZkClient(_connectionConfig, clientConfig);
+      } else {
+        // Use single-zk mode using the ZkAddr given
+        HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();
+        clientConfig.setZkSerializer(new ZNRecordSerializer());
+        zkClient = SharedZkClientFactory.getInstance()
+            .buildZkClient(new HelixZkClient.ZkConnectionConfig(_zkAddr), clientConfig);
+      }
 
       HelixDataAccessor accessor =
           new ZKHelixDataAccessor(_manager.getClusterName(), new ZkBaseDataAccessor<>(zkClient));