public static void addFakeBrokerInstancesToAutoJoinHelixCluster(
     String helixClusterName, String zkServer, int numInstances, boolean isSingleTenant)
     throws Exception {
   for (int i = 0; i < numInstances; ++i) {
     final String brokerId = "Broker_localhost_" + i;
     final HelixManager helixZkManager =
         HelixManagerFactory.getZKHelixManager(
             helixClusterName, brokerId, InstanceType.PARTICIPANT, zkServer);
     final StateMachineEngine stateMachineEngine = helixZkManager.getStateMachineEngine();
     final StateModelFactory<?> stateModelFactory =
         new EmptyBrokerOnlineOfflineStateModelFactory();
     stateMachineEngine.registerStateModelFactory(
         EmptyBrokerOnlineOfflineStateModelFactory.getStateModelDef(), stateModelFactory);
     helixZkManager.connect();
     if (isSingleTenant) {
       helixZkManager
           .getClusterManagmentTool()
           .addInstanceTag(
               helixClusterName,
               brokerId,
               ControllerTenantNameBuilder.getBrokerTenantNameForTenant(
                   ControllerTenantNameBuilder.DEFAULT_TENANT_NAME));
     } else {
       helixZkManager
           .getClusterManagmentTool()
           .addInstanceTag(helixClusterName, brokerId, UNTAGGED_BROKER_INSTANCE);
     }
     Thread.sleep(1000);
   }
 }
  public static void addFakeDataInstancesToAutoJoinHelixCluster(
      String helixClusterName,
      String zkServer,
      int numInstances,
      boolean isSingleTenant,
      int adminPort)
      throws Exception {

    for (int i = 0; i < numInstances; ++i) {
      final String instanceId = "Server_localhost_" + i;

      final HelixManager helixZkManager =
          HelixManagerFactory.getZKHelixManager(
              helixClusterName, instanceId, InstanceType.PARTICIPANT, zkServer);
      final StateMachineEngine stateMachineEngine = helixZkManager.getStateMachineEngine();
      final StateModelFactory<?> stateModelFactory =
          new EmptySegmentOnlineOfflineStateModelFactory();
      stateMachineEngine.registerStateModelFactory(
          EmptySegmentOnlineOfflineStateModelFactory.getStateModelDef(), stateModelFactory);
      helixZkManager.connect();
      if (isSingleTenant) {
        helixZkManager
            .getClusterManagmentTool()
            .addInstanceTag(
                helixClusterName,
                instanceId,
                TableNameBuilder.OFFLINE_TABLE_NAME_BUILDER.forTable(
                    ControllerTenantNameBuilder.DEFAULT_TENANT_NAME));
        helixZkManager
            .getClusterManagmentTool()
            .addInstanceTag(
                helixClusterName,
                instanceId,
                TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(
                    ControllerTenantNameBuilder.DEFAULT_TENANT_NAME));
      } else {
        helixZkManager
            .getClusterManagmentTool()
            .addInstanceTag(helixClusterName, instanceId, UNTAGGED_SERVER_INSTANCE);
      }
      HelixConfigScope scope =
          new HelixConfigScopeBuilder(
                  HelixConfigScope.ConfigScopeProperty.PARTICIPANT, helixClusterName)
              .forParticipant(instanceId)
              .build();
      Map<String, String> props = new HashMap<>();
      props.put(CommonConstants.Helix.Instance.ADMIN_PORT_KEY, String.valueOf(adminPort + i));
      helixZkManager.getClusterManagmentTool().setConfig(scope, props);
    }
  }
  @Test
  public void testZKReconnect() throws Exception {
    final AtomicReference<ZkServer> zkServerRef = new AtomicReference<ZkServer>();
    final int zkPort = TestHelper.getRandomPort();
    final String zkAddr = String.format("localhost:%d", zkPort);
    ZkServer zkServer = TestHelper.startZkServer(zkAddr);
    zkServerRef.set(zkServer);

    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    String clusterName = className + "_" + methodName;

    // Setup cluster
    LOG.info("Setup clusters");
    ClusterSetup clusterSetup = new ClusterSetup(zkAddr);
    clusterSetup.addCluster(clusterName, true);

    // Registers and starts controller
    LOG.info("Starts controller");
    HelixManager controller =
        HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.CONTROLLER, zkAddr);
    controller.connect();

    // Registers and starts participant
    LOG.info("Starts participant");
    String hostname = "localhost";
    String instanceId = String.format("%s_%d", hostname, 1);
    clusterSetup.addInstanceToCluster(clusterName, instanceId);
    HelixManager participant =
        HelixManagerFactory.getZKHelixManager(
            clusterName, instanceId, InstanceType.PARTICIPANT, zkAddr);
    participant.connect();

    LOG.info("Register state machine");
    final CountDownLatch latch = new CountDownLatch(1);
    participant
        .getStateMachineEngine()
        .registerStateModelFactory(
            "OnlineOffline",
            new StateModelFactory<StateModel>() {
              @Override
              public StateModel createNewStateModel(String stateUnitKey) {
                return new SimpleStateModel(latch);
              }
            },
            "test");

    String resourceName = "test-resource";
    LOG.info("Ideal state assignment");
    HelixAdmin helixAdmin = participant.getClusterManagmentTool();
    helixAdmin.addResource(
        clusterName,
        resourceName,
        1,
        "OnlineOffline",
        IdealState.RebalanceMode.CUSTOMIZED.toString());

    IdealState idealState = helixAdmin.getResourceIdealState(clusterName, resourceName);
    idealState.setReplicas("1");
    idealState.setStateModelFactoryName("test");
    idealState.setPartitionState(resourceName + "_0", instanceId, "ONLINE");

    LOG.info("Shutdown ZK server");
    TestHelper.stopZkServer(zkServerRef.get());
    Executors.newSingleThreadScheduledExecutor()
        .schedule(
            new Runnable() {

              @Override
              public void run() {
                try {
                  LOG.info("Restart ZK server");
                  // zkServer.set(TestUtils.startZookeeper(zkDir, zkPort));
                  zkServerRef.set(TestHelper.startZkServer(zkAddr, null, false));
                } catch (Exception e) {
                  LOG.error(e.getMessage(), e);
                }
              }
            },
            2L,
            TimeUnit.SECONDS);

    // future.get();

    LOG.info("Before update ideal state");
    helixAdmin.setResourceIdealState(clusterName, resourceName, idealState);
    LOG.info("After update ideal state");

    LOG.info("Wait for OFFLINE->ONLINE state transition");
    try {
      Assert.assertTrue(latch.await(10, TimeUnit.SECONDS));

      // wait until stable state
      boolean result =
          ClusterStateVerifier.verifyByZkCallback(
              new BestPossAndExtViewZkVerifier(zkAddr, clusterName));
      Assert.assertTrue(result);

    } finally {
      participant.disconnect();
      zkServerRef.get().shutdown();
    }
  }
  public HelixBrokerStarter(
      String helixClusterName, String zkServer, Configuration pinotHelixProperties)
      throws Exception {
    _liveInstancesListener = new LiveInstancesChangeListenerImpl(helixClusterName);

    _pinotHelixProperties = DefaultHelixBrokerConfig.getDefaultBrokerConf(pinotHelixProperties);
    final String brokerId =
        _pinotHelixProperties.getString(
            "instanceId",
            CommonConstants.Helix.PREFIX_OF_BROKER_INSTANCE
                + NetUtil.getHostAddress()
                + "_"
                + _pinotHelixProperties.getInt(
                    CommonConstants.Helix.KEY_OF_BROKER_QUERY_PORT,
                    CommonConstants.Helix.DEFAULT_BROKER_QUERY_PORT));

    _pinotHelixProperties.addProperty("pinot.broker.id", brokerId);
    RoutingTableBuilder defaultOfflineRoutingTableBuilder =
        getRoutingTableBuilder(
            _pinotHelixProperties.subset(DEFAULT_OFFLINE_ROUTING_TABLE_BUILDER_KEY));
    RoutingTableBuilder defaultRealtimeRoutingTableBuilder =
        getRoutingTableBuilder(
            _pinotHelixProperties.subset(DEFAULT_REALTIME_ROUTING_TABLE_BUILDER_KEY));
    Map<String, RoutingTableBuilder> tableToRoutingTableBuilderMap =
        getTableToRoutingTableBuilderMap(_pinotHelixProperties.subset(ROUTING_TABLE_BUILDER_KEY));
    ZkClient zkClient =
        new ZkClient(
            StringUtil.join(
                "/", StringUtils.chomp(zkServer, "/"), helixClusterName, "PROPERTYSTORE"),
            ZkClient.DEFAULT_SESSION_TIMEOUT,
            ZkClient.DEFAULT_CONNECTION_TIMEOUT,
            new ZNRecordSerializer());
    _propertyStore =
        new ZkHelixPropertyStore<ZNRecord>(new ZkBaseDataAccessor<ZNRecord>(zkClient), "/", null);
    _helixExternalViewBasedRouting =
        new HelixExternalViewBasedRouting(
            defaultOfflineRoutingTableBuilder,
            defaultRealtimeRoutingTableBuilder,
            tableToRoutingTableBuilderMap,
            _propertyStore);

    // _brokerServerBuilder = startBroker();
    _brokerServerBuilder = startBroker(_pinotHelixProperties);
    _helixManager =
        HelixManagerFactory.getZKHelixManager(
            helixClusterName, brokerId, InstanceType.PARTICIPANT, zkServer);
    final StateMachineEngine stateMachineEngine = _helixManager.getStateMachineEngine();
    final StateModelFactory<?> stateModelFactory =
        new BrokerResourceOnlineOfflineStateModelFactory(
            _helixManager, _helixExternalViewBasedRouting);
    stateMachineEngine.registerStateModelFactory(
        BrokerResourceOnlineOfflineStateModelFactory.getStateModelDef(), stateModelFactory);
    _helixManager.connect();
    _helixAdmin = _helixManager.getClusterManagmentTool();
    _helixBrokerRoutingTable =
        new HelixBrokerRoutingTable(_helixExternalViewBasedRouting, brokerId, _helixManager);
    addInstanceTagIfNeeded(helixClusterName, brokerId);
    _helixManager.addExternalViewChangeListener(_helixBrokerRoutingTable);
    _helixManager.addInstanceConfigChangeListener(_helixBrokerRoutingTable);
    _helixManager.addLiveInstanceChangeListener(_liveInstancesListener);
  }