org.apache.helix.HelixManagerFactory Java Examples

@BeforeClass
public void beforeClass() throws Exception {
  _participants = new MockParticipantManager[NUM_NODES];
  _gSetupTool.addCluster(CLUSTER_NAME, true);

  _participants = new MockParticipantManager[NUM_NODES];
  for (int i = 0; i < NUM_NODES; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

  for (int i = 0; i < NUM_NODES; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].syncStart();
  }

  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();
}
 

private HelixPartitionSpectator(String zkAddr, String clusterName, String instanceName) {
    logger.info("init HelixPartitionSpectator with zkAddr @{}, clusterName {}, instanceName {}",
            zkAddr, clusterName, instanceName);
    manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
            InstanceType.SPECTATOR, zkAddr);
    zkAddress = zkAddr;
    helixClusterName = clusterName;
    try {
        manager.connect();
        routingTableProvider = new RoutingTableProvider();
        manager.addExternalViewChangeListener(routingTableProvider);


    } catch (Exception e) {
        logger.error("caught exception when init HelixPartitionSpectator", e);
        if (manager != null) {
            manager.disconnect();
        }
        throw new RuntimeException("init HelixPartitionSpectator failure");
    }
}
 

public boolean register(final String serviceId, final ServiceMetadata serviceMetadata)
    throws Exception {
  HelixManager helixManager =
      HelixManagerFactory.getZKHelixManager(cluster, serviceId, InstanceType.PARTICIPANT,
          zkAddress);
  LiveInstanceInfoProvider liveInstanceInfoProvider = new LiveInstanceInfoProvider() {
    @Override
    public ZNRecord getAdditionalLiveInstanceInfo() {
      // serialize serviceMetadata to ZNRecord
      ZNRecord rec = new ZNRecord(serviceId);
      rec.setSimpleField("HOST", serviceMetadata.getHost());
      rec.setSimpleField("PORT", String.valueOf(serviceMetadata.getPort()));
      rec.setSimpleField("SERVICE_NAME", serviceMetadata.getServiceName());
      return rec;
    }
  };
  helixManager.setLiveInstanceInfoProvider(liveInstanceInfoProvider);
  helixManager.connect();
  serviceMap.put(serviceId, helixManager);
  refreshCache();
  return true;
}
 

/**
 * Helper method to start a fake server that only implements Helix part.
 *
 * @throws Exception
 */
private void startFakeServer()
    throws Exception {
  _serverInstance = CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE + NetUtil.getHostAddress() + "_"
      + CommonConstants.Helix.DEFAULT_SERVER_NETTY_PORT;

  // Create server instance with the fake server state model
  _serverHelixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), _serverInstance, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  _serverHelixManager.getStateMachineEngine()
      .registerStateModelFactory(SegmentOnlineOfflineStateModelFactory.getStateModelName(),
          new FakeServerSegmentStateModelFactory());
  _serverHelixManager.connect();

  // Add Helix tag to the server
  _serverHelixManager.getClusterManagmentTool().addInstanceTag(getHelixClusterName(), _serverInstance,
      TableNameBuilder.REALTIME.tableNameWithType(TagNameUtils.DEFAULT_TENANT_NAME));

  // Initialize controller leader locator
  ControllerLeaderLocator.create(_serverHelixManager);
}
 

private void initHelixManager() {
  String zkConnectionString =
      this.clusterConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  logger.info("Using ZooKeeper connection string: " + zkConnectionString);

  if (this.isTaskDriver && this.dedicatedTaskDriverCluster) {
    // This will create a Helix manager to receive the planning job
    this.taskDriverHelixManager = Optional.of(HelixManagerFactory.getZKHelixManager(
        ConfigUtils.getString(this.clusterConfig, GobblinClusterConfigurationKeys.TASK_DRIVER_CLUSTER_NAME_KEY, ""),
        this.helixInstanceName,
        InstanceType.PARTICIPANT,
        zkConnectionString));
    this.jobHelixManager = HelixManagerFactory.getZKHelixManager(
        this.clusterName,
        this.helixInstanceName,
        InstanceType.ADMINISTRATOR,
        zkConnectionString);
  } else {
    this.jobHelixManager = HelixManagerFactory.getZKHelixManager(
        this.clusterName,
        this.helixInstanceName,
        InstanceType.PARTICIPANT,
        zkConnectionString);
  }
}
 

protected void addFakeServerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant, int adminPort)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeSegmentOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeSegmentOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getOfflineTagForTenant(null));
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getRealtimeTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_SERVER_INSTANCE);
  }
  HelixConfigScope configScope =
      new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, getHelixClusterName())
          .forParticipant(instanceId).build();
  helixAdmin.setConfig(configScope, Collections.singletonMap(ADMIN_PORT_KEY, Integer.toString(adminPort)));
  _fakeInstanceHelixManagers.add(helixManager);
}
 

@Test
public void testConnectHelixManagerWithRetry() {
  HelixManager instanceManager = HelixManagerFactory.getZKHelixManager(
      clusterName, corruptHelixInstance, InstanceType.PARTICIPANT, testingZKServer.getConnectString());

  ClusterIntegrationTestUtils.createPartialInstanceStructure(instanceManager, testingZKServer.getConnectString());

  //Ensure that the connecting to Helix without retry will throw a HelixException
  try {
    corruptGobblinTaskRunner.connectHelixManager();
    Assert.fail("Unexpected success in connecting to HelixManager");
  } catch (Exception e) {
    //Assert that a HelixException is thrown.
    Assert.assertTrue(e.getClass().equals(HelixException.class));
  }

  //Ensure that connect with retry succeeds
  corruptGobblinTaskRunner.connectHelixManagerWithRetry();
  Assert.assertTrue(true);
}
 

@Test (groups = {"disabledOnTravis"})
public void testTaskAssignmentAfterHelixConnectionRetry()
    throws Exception {
  Config jobConfigOverrides = ClusterIntegrationTestUtils.buildSleepingJob(JOB_ID, TASK_STATE_FILE);
  this.suite = new TaskAssignmentAfterConnectionRetry(jobConfigOverrides);

  String zkConnectString = suite.getManagerConfig().getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  String clusterName = suite.getManagerConfig().getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);
  //A test manager instance for observing the state of the cluster
  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "TestManager", InstanceType.SPECTATOR, zkConnectString);

  suite.startCluster();

  helixManager.connect();

  //Ensure that Helix has created a workflow
  AssertWithBackoff.create().maxSleepMs(1000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskStarted(helixManager, JOB_ID), "Waiting for the job to start...");

  //Ensure that the SleepingTask is running
  AssertWithBackoff.create().maxSleepMs(100).timeoutMs(2000).backoffFactor(1).
      assertTrue(ClusterIntegrationTest.isTaskRunning(TASK_STATE_FILE),"Waiting for the task to enter running state");

  helixManager.disconnect();
}
 

public void start() throws Exception {
  manager =
      HelixManagerFactory.getZKHelixManager(clusterName, instanceName, InstanceType.PARTICIPANT,
          zkConnectString);

  stateModelFactory = new BootstrapHandler();
  // genericStateMachineHandler = new StateMachineEngine();
  // genericStateMachineHandler.registerStateModelFactory("MasterSlave", stateModelFactory);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory("MasterSlave", stateModelFactory);

  manager.getMessagingService().registerMessageHandlerFactory(
      MessageType.STATE_TRANSITION.name(), stateMach);
  manager.getMessagingService().registerMessageHandlerFactory(
      MessageType.USER_DEFINE_MSG.name(), new CustomMessageHandlerFactory());
  manager.connect();
}
 

public void start() throws Exception {
  manager =
      HelixManagerFactory.getZKHelixManager(clusterName, instanceName, InstanceType.PARTICIPANT,
          zkConnectString);

  if ("MasterSlave".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new MasterSlaveStateModelFactory(this.instanceName, delay);
  } else if ("OnlineOffline".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new OnlineOfflineStateModelFactory(this.instanceName, delay);
  } else if ("LeaderStandby".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new LeaderStandbyStateModelFactory(this.instanceName, delay);
  }
  // genericStateMachineHandler = new StateMachineEngine();
  // genericStateMachineHandler.registerStateModelFactory(stateModelType,
  // stateModelFactory);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory(stateModelType, stateModelFactory);
  manager.connect();
  manager.getMessagingService().registerMessageHandlerFactory(
      MessageType.STATE_TRANSITION.name(), stateMach);
}
 

@Test() public void testInstanceHistory() throws Exception {
  HelixManager manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  manager.connect();

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
  PropertyKey propertyKey = keyBuilder.participantHistory(_participants[0].getInstanceName());
  ParticipantHistory history = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(history);
  List<String> list = history.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 1);

  for (int i = 0; i <= 22; i++) {
    _participants[0].disconnect();
    _participants[0].connect();
  }

  history = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(history);
  list = history.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 20);
  list = history.getRecord().getListField("OFFLINE");
  Assert.assertEquals(list.size(), 20);
  manager.disconnect();
}
 

@Override
public void onBecomeLeaderFromStandby(Message message, NotificationContext context)
    throws Exception {
  String clusterName = message.getPartitionName();
  String controllerName = message.getTgtName();

  logger.info(controllerName + " becoming leader from standby for " + clusterName);

  if (_controller == null) {
    _controller =
        HelixManagerFactory.getZKHelixManager(clusterName, controllerName,
            InstanceType.CONTROLLER, _zkAddr);
    _controller.setEnabledControlPipelineTypes(_enabledPipelineTypes);
    _controller.connect();
    _controller.startTimerTasks();
    logStateTransition("STANDBY", "LEADER", clusterName, controllerName);
  } else {
    logger.error("controller already exists:" + _controller.getInstanceName() + " for "
        + clusterName);
  }

}
 

protected void addFakeBrokerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeBrokerResourceOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeBrokerResourceOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getBrokerTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_BROKER_INSTANCE);
  }
  _fakeInstanceHelixManagers.add(helixManager);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  String storageNodeName = PARTICIPANT_PREFIX + "_" + START_PORT;
  _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  _participant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, storageNodeName);
  _participant.syncStart();

  // create cluster manager
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
}
 

@Test
public void TestSpectator() throws Exception {
  HelixManager relayHelixManager =
      HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, null, InstanceType.SPECTATOR, ZK_ADDR);

  relayHelixManager.connect();
  relayHelixManager.addExternalViewChangeListener(this);

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "NextDB", 64, STATE_MODEL);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "NextDB", 3);

  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));

  Assert.assertTrue(result);

  Assert.assertTrue(_externalViewChanges.containsKey("NextDB"));
  Assert.assertTrue(_externalViewChanges.containsKey(TEST_DB));

}
 

public HelixManager start() throws Exception {
  HelixManager manager = null;
  // zk cluster manager
  if (_clusterMangerType.equalsIgnoreCase("zk")) {
    manager =
        HelixManagerFactory.getZKHelixManager(_clusterName, _instanceName,
            InstanceType.PARTICIPANT, _zkConnectString);
  } else {
    throw new IllegalArgumentException("Unsupported cluster manager type:" + _clusterMangerType);
  }

  MockOnlineOfflineStateModelFactory stateModelFactory2 =
      new MockOnlineOfflineStateModelFactory(_transDelayInMs, _resourceName, _resourceTag,
          _instanceName);
  // genericStateMachineHandler = new StateMachineEngine();
  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory("OnlineOffline", stateModelFactory2);

  manager.connect();
  //manager.getMessagingService().registerMessageHandlerFactory(MessageType.STATE_TRANSITION.name(), genericStateMachineHandler);
  return manager;
}
 

@SuppressWarnings("deprecation")
@Test
public void testSpectator() throws Exception {
  Map<String, Integer> stateReachedCounts = Maps.newHashMap();
  HelixManager participant =
      HelixManagerFactory.getZKHelixManager(_clusterName, "localhost_12918",
          InstanceType.PARTICIPANT, ZK_ADDR);
  participant.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new MyStateModelFactory(stateReachedCounts));
  participant.connect();

  RoutingTableProvider routingTableProvider = new RoutingTableProvider();
  try {
    HelixManager spectator = HelixManagerFactory
        .getZKHelixManager(_clusterName, "spectator", InstanceType.SPECTATOR, ZK_ADDR);
    spectator.connect();
    spectator.addConfigChangeListener(routingTableProvider);
    spectator.addExternalViewChangeListener(routingTableProvider);
    Thread.sleep(1000);

    // Now let's stop the ZK server; this should do nothing
    TestHelper.stopZkServer(_zkServer);
    Thread.sleep(1000);

    // Verify routing table still works
    Assert.assertEquals(routingTableProvider.getInstances("resource0", "ONLINE").size(), 1);
    Assert.assertEquals(routingTableProvider.getInstances("resource0", "OFFLINE").size(), 0);
  } finally {
    routingTableProvider.shutdown();
    if (participant.isConnected()) {
      participant.disconnect();
    }
  }
}
 

@Test
public void testParticipant() throws Exception {
  Map<String, Integer> stateReachedCounts = Maps.newHashMap();
  HelixManager participant =
      HelixManagerFactory.getZKHelixManager(_clusterName, "localhost_12918",
          InstanceType.PARTICIPANT, ZK_ADDR);
  participant.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new MyStateModelFactory(stateReachedCounts));
  participant.connect();

  Thread.sleep(1000);

  // Ensure that the external view coalesces
  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          _clusterName));
  Assert.assertTrue(result);

  // Ensure that there was only one state transition
  Assert.assertEquals(stateReachedCounts.size(), 1);
  Assert.assertTrue(stateReachedCounts.containsKey("ONLINE"));
  Assert.assertEquals(stateReachedCounts.get("ONLINE").intValue(), 1);

  // Now let's stop the ZK server; this should do nothing
  TestHelper.stopZkServer(_zkServer);
  Thread.sleep(1000);

  // Verify no change
  Assert.assertEquals(stateReachedCounts.size(), 1);
  Assert.assertTrue(stateReachedCounts.containsKey("ONLINE"));
  Assert.assertEquals(stateReachedCounts.get("ONLINE").intValue(), 1);

  if (participant.isConnected()) {
    participant.disconnect();
  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _numPartitions = 1;
  _numNodes = 3;
  super.beforeClass();
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);

  // Stop participants that have been started in super class
  for (int i = 0; i < _numNodes; i++) {
    super.stopParticipant(i);
    Assert.assertFalse(_participants[i].isConnected());
  }

  // Start new participants that have new TaskStateModel (NewMockTask) information
  _participants = new MockParticipantManager[_numNodes];
  for (int i = 0; i < _numNodes; i++) {
    Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
    taskFactoryReg.put(NewMockTask.TASK_COMMAND, NewMockTask::new);
    String instanceName = PARTICIPANT_PREFIX + "_" + (_startPort + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);

    // Register a Task state model factory.
    StateMachineEngine stateMachine = _participants[i].getStateMachineEngine();
    stateMachine.registerStateModelFactory("Task",
        new TaskStateModelFactory(_participants[i], taskFactoryReg));
    _participants[i].syncStart();
  }

  _manager.connect();
  _driver = new TaskDriver(_manager);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  super.beforeClass();
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, _PARTITIONS, STATE_MODEL);
  for (int i = 0; i < NODE_NR; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB, _replica);

  // start dummy participants
  for (int i = 0; i < NODE_NR; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].syncStart();
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  _clusterVerifier = new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // create cluster manager
  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "Admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
}
 

public void initialize() throws Exception {
  manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
    InstanceType.PARTICIPANT, zkAddress);

  OnlineOfflineStateModelFactory stateModelFactory =
    new OnlineOfflineStateModelFactory(instanceName, this);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory(MetricCollectorHAController.DEFAULT_STATE_MODEL, stateModelFactory);
  manager.connect();

  checkpointManager = new CheckpointManager(manager.getHelixPropertyStore());
}
 

@Override
@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);

  for (int i = 0; i < NODE_NR; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "Admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _configAccessor = new ConfigAccessor(_gZkClient);

  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller =
      new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  boolean result =
      ClusterStateVerifier
          .verifyByZkCallback(new MasterNbInExtViewVerifier(ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);
}
 

private HelixManager initializeHelixManager(String clusterName, String instanceName) {
  HelixManager manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
      InstanceType.PARTICIPANT, org.apache.helix.common.ZkTestBase.ZK_ADDR);

  MasterSlaveStateModelFactory stateModelFactory = new MasterSlaveStateModelFactory(instanceName);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory("id1", stateModelFactory);
  return manager;
}
 

@BeforeClass
public void beforeClass()
    throws Exception {
  super.beforeClass();

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, _PARTITIONS, STATE_MODEL);
  for (int i = 0; i < NODE_NR; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB, _replica);

  // start dummy participants
  for (int i = 0; i < NODE_NR; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].syncStart();
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  // create cluster manager
  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "Admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

  // create AssignmentMetadataStore
  _store = new AssignmentMetadataStore(_manager.getMetadataStoreConnectionString(),
      _manager.getClusterName());
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "Test", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _configAccessor = new ConfigAccessor(_gZkClient);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  super.beforeClass();
  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _participants[0].connect();
  _mockProvider = CustomizedStateProviderFactory.getInstance()
      .buildCustomizedStateProvider(_manager, _participants[0].getInstanceName());
  _dataAccessor = _manager.getHelixDataAccessor();
  _propertyKey = _dataAccessor.keyBuilder()
      .customizedStates(_participants[0].getInstanceName(), CUSTOMIZE_STATE_NAME);
}
 

public PinotZKChanger(String zkAddress, String clusterName) {
  this.clusterName = clusterName;
  helixAdmin = new ZKHelixAdmin(zkAddress);
  helixManager = HelixManagerFactory
      .getZKHelixManager(clusterName, "PinotNumReplicaChanger", InstanceType.ADMINISTRATOR, zkAddress);
  try {
    helixManager.connect();
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  ZNRecordSerializer serializer = new ZNRecordSerializer();
  String path = PropertyPathConfig.getPath(PropertyType.PROPERTYSTORE, clusterName);
  propertyStore = new ZkHelixPropertyStore<>(zkAddress, serializer, path);
}
 

/**
 * @return returns true if
 */
public boolean start() throws Exception {
  // auto create cluster and allow nodes to automatically join the cluster
  admin =
      HelixManagerFactory.getZKHelixManager(cluster, "service-discovery",
          InstanceType.ADMINISTRATOR, zkAddress);
  admin.connect();
  admin.getClusterManagmentTool().addCluster(cluster, false);
  HelixConfigScope scope =
      new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER).forCluster(cluster).build();

  Map<String, String> properties = new HashMap<String, String>();
  properties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
  admin.getClusterManagmentTool().setConfig(scope, properties);
  switch (mode) {
  case POLL:
    startBackgroundTask();
    break;
  case WATCH:
    setupWatcher();
    break;
  case NONE:// dont monitor changes, supports only registration

  }
  refreshCache();
  return true;
}
 

public void start() throws Exception {
  System.out.println("STARTING " + instanceName);
  configureInstance(instanceName);
  participantManager =
      HelixManagerFactory.getZKHelixManager(clusterName, instanceName, InstanceType.PARTICIPANT,
          zkAddress);
  participantManager.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new LockFactory());
  participantManager.connect();
  if (startController) {
    startController();
  }
  System.out.println("STARTED " + instanceName);

}
 

@BeforeClass
public void beforeClass() throws Exception {
  super.beforeClass();

  // Stop participants that have been started in super class
  for (int i = 0; i < _numNodes; i++) {
    super.stopParticipant(i);
  }

  // Check that participants are actually stopped
  for (int i = 0; i < _numNodes; i++) {
    Assert.assertFalse(_participants[i].isConnected());
  }

  // Start new participants that have new TaskStateModel (DelayedStopTask) information
  _participants = new MockParticipantManager[_numNodes];
  for (int i = 0; i < _numNodes; i++) {
    Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
    taskFactoryReg.put(DelayedStopTask.TASK_COMMAND, DelayedStopTask::new);
    String instanceName = PARTICIPANT_PREFIX + "_" + (_startPort + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);

    // Register a Task state model factory.
    StateMachineEngine stateMachine = _participants[i].getStateMachineEngine();
    stateMachine.registerStateModelFactory("Task",
        new TaskStateModelFactory(_participants[i], taskFactoryReg));
    _participants[i].syncStart();
  }

  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

  _driver = new TaskDriver(_manager);

  _admin = _gSetupTool.getClusterManagementTool();
}