Java Code Examples for org.apache.helix.HelixAdmin#addInstance()

@PUT
public Response addInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName, String content) {
  HelixAdmin admin = getHelixAdmin();
  ZNRecord record;
  try {
    record = toZNRecord(content);
  } catch (IOException e) {
    LOG.error("Failed to deserialize user's input " + content + ", Exception: " + e);
    return badRequest("Input is not a vaild ZNRecord!");
  }

  try {
    admin.addInstance(clusterId, new InstanceConfig(record));
  } catch (Exception ex) {
    LOG.error("Error in adding an instance: " + instanceName, ex);
    return serverError(ex);
  }

  return OK();
}
 

@Test
public void testEnableDisablePartitions() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String instanceName = "TestInstance";
  String testResourcePrefix = "TestResource";
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addInstance(clusterName, new InstanceConfig(instanceName));

  // Test disable instances with resources
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "0",
      Arrays.asList("1", "2"));
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "1",
      Arrays.asList("2", "3", "4"));
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 2);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 3);

  // Test disable partition across resources
  // TODO : Remove this part once setInstanceEnabledForPartition(partition, enabled) is removed
  instanceConfig.setInstanceEnabledForPartition("10", false);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 3);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 4);
  admin.dropCluster(clusterName);
}
 

protected void setupInstances(String clusterName, int[] instances) {
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  for (int i = 0; i < instances.length; i++) {
    String instance = "localhost_" + instances[i];
    InstanceConfig instanceConfig = new InstanceConfig(instance);
    instanceConfig.setHostName("localhost");
    instanceConfig.setPort("" + instances[i]);
    instanceConfig.setInstanceEnabled(true);
    admin.addInstance(clusterName, instanceConfig);
  }
}
 

@Test
public void testRebuildBrokerResourceWhenBrokerAdded()
    throws Exception {
  // Check that the first table we added doesn't need to be rebuilt(case where ideal state brokers and brokers in broker resource are the same.
  String partitionName = _offlineTableConfig.getTableName();
  HelixAdmin helixAdmin = _helixManager.getClusterManagmentTool();

  IdealState idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Ensure that the broker resource is not rebuilt.
  Assert.assertTrue(idealState.getInstanceSet(partitionName)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
  _helixResourceManager.rebuildBrokerResourceFromHelixTags(partitionName);

  // Add another table that needs to be rebuilt
  TableConfig offlineTableConfigTwo = new TableConfigBuilder(TableType.OFFLINE).setTableName(TEST_TABLE_TWO).build();
  _helixResourceManager.addTable(offlineTableConfigTwo);
  String partitionNameTwo = offlineTableConfigTwo.getTableName();

  // Add a new broker manually such that the ideal state is not updated and ensure that rebuild broker resource is called
  final String brokerId = "Broker_localhost_2";
  InstanceConfig instanceConfig = new InstanceConfig(brokerId);
  instanceConfig.setInstanceEnabled(true);
  instanceConfig.setHostName("Broker_localhost");
  instanceConfig.setPort("2");
  helixAdmin.addInstance(getHelixClusterName(), instanceConfig);
  helixAdmin.addInstanceTag(getHelixClusterName(), instanceConfig.getInstanceName(),
      TagNameUtils.getBrokerTagForTenant(TagNameUtils.DEFAULT_TENANT_NAME));
  idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Assert that the two don't equal before the call to rebuild the broker resource.
  Assert.assertTrue(!idealState.getInstanceSet(partitionNameTwo)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
  _helixResourceManager.rebuildBrokerResourceFromHelixTags(partitionNameTwo);
  idealState = HelixHelper.getBrokerIdealStates(helixAdmin, getHelixClusterName());
  // Assert that the two do equal after being rebuilt.
  Assert.assertTrue(idealState.getInstanceSet(partitionNameTwo)
      .equals(_helixResourceManager.getAllInstancesForBrokerTenant(TagNameUtils.DEFAULT_TENANT_NAME)));
}
 

/**
 * Test success case when updating InstanceConfig in Helix after new replica is added in storage manager.
 */
@Test
public void updateInstanceConfigSuccessTest() throws Exception {
  generateConfigs(true, true);
  MockDataNodeId localNode = clusterMap.getDataNodes().get(0);
  List<ReplicaId> localReplicas = clusterMap.getReplicaIds(localNode);
  MockClusterParticipant mockHelixParticipant = new MockClusterParticipant();
  StorageManager storageManager =
      createStorageManager(localNode, metricRegistry, Collections.singletonList(mockHelixParticipant));
  storageManager.start();
  // create a new partition and get its replica on local node
  PartitionId newPartition = clusterMap.createNewPartition(Collections.singletonList(localNode));
  ReplicaId newReplica = newPartition.getReplicaIds().get(0);
  // for updating instanceConfig test, we first add an empty InstanceConfig of current node
  String instanceName =
      ClusterMapUtils.getInstanceName(clusterMapConfig.clusterMapHostName, clusterMapConfig.clusterMapPort);
  InstanceConfig instanceConfig = new InstanceConfig(instanceName);
  instanceConfig.setHostName(localNode.getHostname());
  instanceConfig.setPort(Integer.toString(localNode.getPort()));
  // for current test, we initial InstanceConfig empty, non-empty case will be tested in HelixParticipantTest
  Map<String, Map<String, String>> diskInfos = new HashMap<>();
  instanceConfig.getRecord().setMapFields(diskInfos);
  HelixAdmin helixAdmin = mockHelixParticipant.getHelixAdmin();
  helixAdmin.addCluster(CLUSTER_NAME);
  helixAdmin.addInstance(CLUSTER_NAME, instanceConfig);
  // test success case
  mockHelixParticipant.onPartitionBecomeBootstrapFromOffline(newPartition.toPathString());
  instanceConfig = helixAdmin.getInstanceConfig(CLUSTER_NAME, instanceName);
  // verify that new replica info is present in InstanceConfig
  Map<String, Map<String, String>> mountPathToDiskInfos = instanceConfig.getRecord().getMapFields();
  Map<String, String> diskInfo = mountPathToDiskInfos.get(newReplica.getMountPath());
  String replicasStr = diskInfo.get("Replicas");
  Set<String> partitionStrs = new HashSet<>();
  for (String replicaInfo : replicasStr.split(",")) {
    String[] infos = replicaInfo.split(":");
    partitionStrs.add(infos[0]);
  }
  assertTrue("New replica info is not found in InstanceConfig", partitionStrs.contains(newPartition.toPathString()));
  shutdownAndAssertStoresInaccessible(storageManager, localReplicas);
}
 

/**
 * Test Helix Participant creation and addition.
 * Helix Java APIs tested in this method are:
 * ZkHelixAdmin and ZKHelixManager (mock participant/controller)
 */
@Test(dependsOnMethods = "testCreateClusters")
public void testCreateParticipants() throws Exception {
  // Create two ClusterSetups using two different constructors
  // Note: ZK Address here could be anything because multiZk mode is on (it will be ignored)
  HelixAdmin helixAdminZkAddr = new ZKHelixAdmin(ZK_SERVER_MAP.keySet().iterator().next());
  HelixAdmin helixAdminBuilder = new ZKHelixAdmin.Builder().build();
  _zkHelixAdmin = helixAdminBuilder;

  String participantNamePrefix = "Node_";
  int numParticipants = 5;
  createParticipantsAndVerify(helixAdminZkAddr, numParticipants, participantNamePrefix);
  createParticipantsAndVerify(helixAdminBuilder, numParticipants, participantNamePrefix);

  // Create mock controller and participants for next tests
  for (String cluster : CLUSTER_LIST) {
    // Start a controller
    // Note: in multiZK mode, ZK Addr is ignored
    ClusterControllerManager mockController =
        new ClusterControllerManager("DummyZK", cluster, "controller");
    mockController.syncStart();
    MOCK_CONTROLLERS.put(cluster, mockController);

    for (int i = 0; i < numParticipants; i++) {
      // Note: in multiZK mode, ZK Addr is ignored
      InstanceConfig instanceConfig = new InstanceConfig(participantNamePrefix + i);
      helixAdminBuilder.addInstance(cluster, instanceConfig);
      MockParticipantManager mockNode =
          new MockParticipantManager("DummyZK", cluster, participantNamePrefix + i);

      // Register task state model for task framework testing in later methods
      Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
      taskFactoryReg.put(MockTask.TASK_COMMAND, MockTask::new);
      // Register a Task state model factory.
      StateMachineEngine stateMachine = mockNode.getStateMachineEngine();
      stateMachine
          .registerStateModelFactory("Task", new TaskStateModelFactory(mockNode, taskFactoryReg));

      mockNode.syncStart();
      MOCK_PARTICIPANTS.add(mockNode);
    }
    // Check that mockNodes are up
    Assert.assertTrue(TestHelper
        .verify(() -> helixAdminBuilder.getInstancesInCluster(cluster).size() == numParticipants,
            TestHelper.WAIT_DURATION));
  }
}
 

@Test
public void testGetResourcesWithTag() {
  String TEST_TAG = "TestTAG";

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

  HelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient));

  tool.addStateModelDef(clusterName, "OnlineOffline",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline()));

  for (int i = 0; i < 4; i++) {
    String instanceName = "host" + i + "_9999";
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("host" + i);
    config.setPort("9999");
    // set tag to two instances
    if (i < 2) {
      config.addTag(TEST_TAG);
    }
    tool.addInstance(clusterName, config);
    tool.enableInstance(clusterName, instanceName, true);
    String path = PropertyPathBuilder.instance(clusterName, instanceName);
    AssertJUnit.assertTrue(_gZkClient.exists(path));
  }

  for (int i = 0; i < 4; i++) {
    String resourceName = "database_" + i;
    IdealState is = new IdealState(resourceName);
    is.setStateModelDefRef("OnlineOffline");
    is.setNumPartitions(2);
    is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
    is.setReplicas("1");
    is.enable(true);
    if (i < 2) {
      is.setInstanceGroupTag(TEST_TAG);
    }
    tool.addResource(clusterName, resourceName, is);
  }

  List<String> allResources = tool.getResourcesInCluster(clusterName);
  List<String> resourcesWithTag = tool.getResourcesInClusterWithTag(clusterName, TEST_TAG);
  AssertJUnit.assertEquals(allResources.size(), 4);
  AssertJUnit.assertEquals(resourcesWithTag.size(), 2);

  tool.dropCluster(clusterName);
}
 

/**
 * Test addResourceWithWeight() and validateResourcesForWagedRebalance() by trying to add a resource with incomplete ResourceConfig.
 */
@Test
public void testAddResourceWithWeightAndValidation()
    throws IOException {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String mockInstance = "MockInstance";
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Create a dummy instance
  InstanceConfig instanceConfig = new InstanceConfig(mockInstance);
  Map<String, Integer> mockInstanceCapacity =
      ImmutableMap.of("WCU", 100, "RCU", 100, "STORAGE", 100);
  instanceConfig.setInstanceCapacityMap(mockInstanceCapacity);
  admin.addInstance(clusterName, instanceConfig);
  MockParticipantManager mockParticipantManager =
      new MockParticipantManager(ZK_ADDR, clusterName, mockInstance);
  mockParticipantManager.syncStart();

  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);

  ResourceConfig resourceConfig = new ResourceConfig(testResourcePrefix);
  // validate
  Map<String, Boolean> validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));
  try {
    admin.addResourceWithWeight(clusterName, idealState, resourceConfig);
    Assert.fail();
  } catch (HelixException e) {
    // OK since resourceConfig is empty
  }

  // Set PARTITION_CAPACITY_MAP
  Map<String, String> capacityDataMap =
      ImmutableMap.of("WCU", "1", "RCU", "2", "STORAGE", "3");
  resourceConfig.getRecord()
      .setMapField(ResourceConfig.ResourceConfigProperty.PARTITION_CAPACITY_MAP.name(),
          Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY,
              OBJECT_MAPPER.writeValueAsString(capacityDataMap)));

  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));

  // Add the capacity key to ClusterConfig
  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();
  ClusterConfig clusterConfig = dataAccessor.getProperty(keyBuilder.clusterConfig());
  clusterConfig.setInstanceCapacityKeys(Arrays.asList("WCU", "RCU", "STORAGE"));
  dataAccessor.setProperty(keyBuilder.clusterConfig(), clusterConfig);

  // Should succeed now
  Assert.assertTrue(admin.addResourceWithWeight(clusterName, idealState, resourceConfig));
  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertTrue(validationResult.get(testResourcePrefix));
}
 

@Test
public void testGetInstances() throws IOException {
  final String clusterName = "TestTagAwareness_testGetResources";
  final String[] TAGS = {
      "tag1", "tag2"
  };
  final String URL_BASE =
      "http://localhost:" + ADMIN_PORT + "/clusters/" + clusterName + "/instances";

  _gSetupTool.addCluster(clusterName, true);
  HelixAdmin admin = _gSetupTool.getClusterManagementTool();

  // Add 4 participants, each with differint tag characteristics
  InstanceConfig instance1 = new InstanceConfig("localhost_1");
  instance1.addTag(TAGS[0]);
  admin.addInstance(clusterName, instance1);
  InstanceConfig instance2 = new InstanceConfig("localhost_2");
  instance2.addTag(TAGS[1]);
  admin.addInstance(clusterName, instance2);
  InstanceConfig instance3 = new InstanceConfig("localhost_3");
  instance3.addTag(TAGS[0]);
  instance3.addTag(TAGS[1]);
  admin.addInstance(clusterName, instance3);
  InstanceConfig instance4 = new InstanceConfig("localhost_4");
  admin.addInstance(clusterName, instance4);

  // Now make a REST call for all resources
  Reference resourceRef = new Reference(URL_BASE);
  Request request = new Request(Method.GET, resourceRef);
  Response response = _gClient.handle(request);
  ListInstancesWrapper responseWrapper =
      ClusterRepresentationUtil.JsonToObject(ListInstancesWrapper.class,
          response.getEntityAsText());
  Map<String, List<String>> tagInfo = responseWrapper.tagInfo;

  // Ensure tag ownership is reported correctly
  Assert.assertTrue(tagInfo.containsKey(TAGS[0]));
  Assert.assertTrue(tagInfo.containsKey(TAGS[1]));
  Assert.assertTrue(tagInfo.get(TAGS[0]).contains("localhost_1"));
  Assert.assertFalse(tagInfo.get(TAGS[0]).contains("localhost_2"));
  Assert.assertTrue(tagInfo.get(TAGS[0]).contains("localhost_3"));
  Assert.assertFalse(tagInfo.get(TAGS[0]).contains("localhost_4"));
  Assert.assertFalse(tagInfo.get(TAGS[1]).contains("localhost_1"));
  Assert.assertTrue(tagInfo.get(TAGS[1]).contains("localhost_2"));
  Assert.assertTrue(tagInfo.get(TAGS[1]).contains("localhost_3"));
  Assert.assertFalse(tagInfo.get(TAGS[1]).contains("localhost_4"));
}
 

/**
 * Test both replica info addition and removal cases when updating node info in Helix cluster.
 * @throws Exception
 */
@Test
public void testUpdateNodeInfoInCluster() throws Exception {
  // test setup: 3 disks on local node, each disk has 3 replicas
  MockClusterMap clusterMap = new MockClusterMap(false, true, 1, 3, 3, false, false);
  MockDataNodeId localNode = clusterMap.getDataNodes().get(0);
  List<ReplicaId> localReplicas = clusterMap.getReplicaIds(localNode);
  ReplicaId existingReplica = localReplicas.get(0);
  // override some props for current test
  props.setProperty("clustermap.update.datanode.info", Boolean.toString(true));
  props.setProperty("clustermap.port", String.valueOf(localNode.getPort()));
  ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(props));
  HelixParticipant participant = new HelixParticipant(clusterMapConfig, helixManagerFactory, new MetricRegistry(),
      getDefaultZkConnectStr(clusterMapConfig), true);
  // create InstanceConfig for local node. Also, put existing replica into sealed list
  List<String> sealedList = new ArrayList<>();
  sealedList.add(existingReplica.getPartitionId().toPathString());
  InstanceConfig instanceConfig = generateInstanceConfig(clusterMap, localNode, sealedList);
  HelixAdmin helixAdmin = participant.getHelixAdmin();
  helixAdmin.addCluster(clusterMapConfig.clusterMapClusterName);
  helixAdmin.addInstance(clusterMapConfig.clusterMapClusterName, instanceConfig);
  String instanceName = ClusterMapUtils.getInstanceName(localNode.getHostname(), localNode.getPort());
  // generate exactly same config for comparison
  InstanceConfig initialInstanceConfig = generateInstanceConfig(clusterMap, localNode, sealedList);
  // 1. add existing replica's info to Helix should be no-op
  assertTrue("Adding existing replica's info should succeed",
      participant.updateDataNodeInfoInCluster(existingReplica, true));
  assertEquals("InstanceConfig should stay unchanged", initialInstanceConfig,
      helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName));
  // create two new partitions on the same disk of local node
  PartitionId newPartition1 = clusterMap.createNewPartition(Collections.singletonList(localNode), 0);
  PartitionId newPartition2 = clusterMap.createNewPartition(Collections.singletonList(localNode), 0);
  // 2. add new partition2 (id = 10, replicaFromPartition2) into InstanceConfig
  ReplicaId replicaFromPartition2 = newPartition2.getReplicaIds().get(0);
  assertTrue("Adding new replica info into InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition2, true));
  // verify new added replica (id = 10, replicaFromPartition2) info is present in InstanceConfig
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 3. add new partition1 (id = 9, replicaFromPartition1) into InstanceConfig
  ReplicaId replicaFromPartition1 = newPartition1.getReplicaIds().get(0);
  assertTrue("Adding new replica info into InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, true));
  // verify new added replica (id = 9, replicaFromPartition1) info is present in InstanceConfig
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition1, true);
  // ensure previous added replica (id = 10, replicaFromPartition2) still exists
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 4. remove recently added new replica (id = 9, replicaFromPartition1)
  assertTrue("Removing replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, false));
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition1, false);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // 5. remove same replica again (id = 9, replicaFromPartition1) should be no-op
  assertTrue("Removing non-found replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(replicaFromPartition1, false));
  // 6. remove an existing replica should succeed
  assertTrue("Removing replica info from InstanceConfig should succeed.",
      participant.updateDataNodeInfoInCluster(existingReplica, false));
  instanceConfig = helixAdmin.getInstanceConfig(clusterMapConfig.clusterMapClusterName, instanceName);
  verifyReplicaInfoInInstanceConfig(instanceConfig, existingReplica, false);
  verifyReplicaInfoInInstanceConfig(instanceConfig, replicaFromPartition2, true);
  // reset props
  props.setProperty("clustermap.update.datanode.info", Boolean.toString(false));
  props.setProperty("clustermap.port", "2200");
}