@Test
  public void simpleIntegrationTest() throws Exception {
    // Logger.getRootLogger().setLevel(Level.INFO);
    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    String clusterName = className + "_" + methodName;
    int n = 1;

    System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

    TestHelper.setupCluster(
        clusterName,
        _zkaddr,
        12918, // participant port
        "localhost", // participant name prefix
        "TestDB", // resource name prefix
        1, // resources
        4, // partitions per resource
        n, // number of nodes
        1, // replicas
        "MasterSlave",
        true); // do rebalance

    HelixManager participant =
        new ZKHelixManager(clusterName, "localhost_12918", InstanceType.PARTICIPANT, _zkaddr);
    participant
        .getStateMachineEngine()
        .registerStateModelFactory(StateModelDefId.MasterSlave, new MockMSModelFactory());
    participant.connect();

    HelixManager controller =
        new ZKHelixManager(clusterName, "controller_0", InstanceType.CONTROLLER, _zkaddr);
    controller.connect();

    boolean result =
        ClusterStateVerifier.verifyByZkCallback(
            new BestPossAndExtViewZkVerifier(_zkaddr, clusterName));
    Assert.assertTrue(result);

    // cleanup
    controller.disconnect();
    participant.disconnect();

    // verify all live-instances and leader nodes are gone
    ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    Assert.assertNull(accessor.getProperty(keyBuilder.liveInstance("localhost_12918")));
    Assert.assertNull(accessor.getProperty(keyBuilder.controllerLeader()));

    System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
  }
 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);
    }
  }
 void start() {
   try {
     participantManager
         .getMessagingService()
         .registerMessageHandlerFactory(
             Message.MessageType.USER_DEFINE_MSG.toString(),
             new MessageHandlerResolverWrapper().convert());
     participantManager
         .getStateMachineEngine()
         .registerStateModelFactory("LeaderStandby", new LockTransitionalFactory());
     participantManager.connect();
     offlinePartition();
   } catch (final Exception ex) {
     throw new RuntimeException(ex);
   }
 }
  @Test
  public void testBounceAll() throws Exception {
    // pick numbers that don't divide evenly
    final int NUM_PARTICIPANTS = 5;
    final int NUM_PARTITIONS = 123;
    final int NUM_REPLICAS = 1;
    final String RESOURCE_PREFIX = "TestDB";
    final String RESOURCE_NAME = RESOURCE_PREFIX + "0";

    // create a cluster name based on this test name
    String className = TestHelper.getTestClassName();
    String methodName = TestHelper.getTestMethodName();
    String clusterName = className + "_" + methodName;
    System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

    // Set up cluster
    TestHelper.setupCluster(
        clusterName,
        ZK_ADDR,
        12918, // participant port
        "localhost", // participant name prefix
        "TestDB", // resource name prefix
        1, // resources
        NUM_PARTITIONS, // partitions per resource
        NUM_PARTICIPANTS, // number of nodes
        NUM_REPLICAS, // replicas
        "OnlineOffline",
        RebalanceMode.FULL_AUTO, // use FULL_AUTO mode to test node tagging
        true); // do rebalance

    // Start the participants
    HelixManager[] participants = new HelixManager[NUM_PARTICIPANTS];
    for (int i = 0; i < NUM_PARTICIPANTS; i++) {
      final String instanceName = "localhost_" + (12918 + i);
      participants[i] = createParticipant(clusterName, instanceName);
      participants[i].connect();
    }

    // Start the controller
    ClusterControllerManager controller =
        new ClusterControllerManager(ZK_ADDR, clusterName, "controller");
    controller.syncStart();

    // get an admin and accessor
    HelixAdmin helixAdmin = new ZKHelixAdmin(_gZkClient);
    BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
    HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();

    // do the test
    try {
      Thread.sleep(1000);
      // ensure that the external view coalesces
      boolean result =
          ClusterStateVerifier.verifyByZkCallback(
              new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
      Assert.assertTrue(result);
      ExternalView stableExternalView =
          accessor.getProperty(keyBuilder.externalView(RESOURCE_NAME));
      for (HelixManager participant : participants) {
        // disable the controller, bounce the node, re-enable the controller, verify assignments
        // remained the same
        helixAdmin.enableCluster(clusterName, false);
        participant.disconnect();
        Thread.sleep(1000);
        participant = createParticipant(clusterName, participant.getInstanceName());
        participant.connect();
        Thread.sleep(1000);
        helixAdmin.enableCluster(clusterName, true);
        Thread.sleep(1000);
        result =
            ClusterStateVerifier.verifyByZkCallback(
                new MatchingExternalViewVerifier(stableExternalView, clusterName));
        Assert.assertTrue(result);
      }
    } finally {
      // clean up
      controller.syncStop();
      for (HelixManager participant : participants) {
        participant.disconnect();
      }
      System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
    }
  }
  @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);
  }