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

@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);
}
 

@DELETE
public Response deleteInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName) {
  HelixAdmin admin = getHelixAdmin();
  try {
    InstanceConfig instanceConfig = admin.getInstanceConfig(clusterId, instanceName);
    admin.dropInstance(clusterId, instanceConfig);
  } catch (HelixException e) {
    return badRequest(e.getMessage());
  }

  return OK();
}
 

/**
 * Tags Pinot Minion instance if needed.
 */
private void addInstanceTagIfNeeded() {
  HelixAdmin helixAdmin = _helixManager.getClusterManagmentTool();
  String clusterName = _helixManager.getClusterName();
  InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, _instanceId);
  if (instanceConfig.getTags().isEmpty()) {
    LOGGER.info("Adding default Helix tag: {} to Pinot minion", CommonConstants.Helix.UNTAGGED_MINION_INSTANCE);
    helixAdmin.addInstanceTag(clusterName, _instanceId, CommonConstants.Helix.UNTAGGED_MINION_INSTANCE);
  }
}
 

/**
 * Get sealed partitions from given datacenter.
 * @param dc the datacenter where sealed partitions come from.
 * @param dcToSealedPartitions a map to track sealed partitions in each dc. This entry associated with given dc will
 *                             be populated in this method.
 * @param nodeToNonExistentReplicas a map to track if any replica is in sealed list but not actually on local node.
 */
private void getSealedPartitionsInDc(Datacenter dc, Map<String, Set<String>> dcToSealedPartitions,
    Map<String, Set<String>> nodeToNonExistentReplicas) {
  String dcName = dc.getName();
  dcToSealedPartitions.put(dcName, new HashSet<>());
  HelixAdmin admin = adminForDc.get(dcName);
  Set<String> allInstancesInHelix = new HashSet<>(admin.getInstancesInCluster(clusterName));
  for (DataNodeId dataNodeId : dc.getDataNodes()) {
    DataNode dataNode = (DataNode) dataNodeId;
    Set<String> replicasOnNode = staticClusterMap.getReplicas(dataNode)
        .stream()
        .map(replicaId -> replicaId.getPartitionId().toPathString())
        .collect(Collectors.toSet());
    String instanceName = getInstanceName(dataNode);
    ensureOrThrow(allInstancesInHelix.contains(instanceName), "Instance not present in Helix " + instanceName);
    InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
    List<String> sealedReplicas = instanceConfig.getRecord().getListField(ClusterMapUtils.SEALED_STR);
    if (sealedReplicas != null) {
      for (String sealedReplica : sealedReplicas) {
        info("Replica {} is sealed on {}", sealedReplica, instanceName);
        dcToSealedPartitions.get(dcName).add(sealedReplica);
        if (!replicasOnNode.contains(sealedReplica)) {
          logger.warn("Replica {} is in sealed list but not on node {}", sealedReplica, instanceName);
          nodeToNonExistentReplicas.computeIfAbsent(instanceName, key -> new HashSet<>()).add(sealedReplica);
        }
      }
    }
  }
}
 

/**
 * 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);
}
 

/**
 * Control state of partition on certain node. (i.e. DisablePartition, EnablePartition)
 */
private void controlPartitionState() {
  // for now, we only support controlling state of single partition on certain node. Hence, there should be only 1 dc
  // in adminForDc map.
  if (adminForDc.size() != 1) {
    throw new IllegalStateException("The dc count is not 1 for partition state control operation");
  }
  HelixAdmin helixAdmin = adminForDc.values().iterator().next();
  String instanceName;
  if (portNum == null) {
    Optional<DataNodeId> optionalDataNode =
        staticClusterMap.getDataNodeIds().stream().filter(node -> node.getHostname().equals(hostName)).findFirst();
    if (!optionalDataNode.isPresent()) {
      throw new IllegalStateException("Host " + hostName + " is not found in static clustermap");
    }
    DataNode dataNode = (DataNode) optionalDataNode.get();
    instanceName = getInstanceName(dataNode);
  } else {
    instanceName = ClusterMapUtils.getInstanceName(hostName, portNum);
  }
  InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
  String resourceNameForPartition = getResourceNameOfPartition(helixAdmin, clusterName, partitionName);
  info("{} partition {} under resource {} on node {}",
      helixAdminOperation == HelixAdminOperation.EnablePartition ? "Enabling" : "Disabling", partitionName,
      resourceNameForPartition, instanceName);
  instanceConfig.setInstanceEnabledForPartition(resourceNameForPartition, partitionName,
      helixAdminOperation == HelixAdminOperation.EnablePartition);
  // clean up the disabled partition entry if it exists and is empty.
  Map<String, String> disabledPartitions =
      instanceConfig.getRecord().getMapFields().get(HELIX_DISABLED_PARTITION_STR);
  if (disabledPartitions != null && disabledPartitions.isEmpty()) {
    instanceConfig.getRecord().getMapFields().remove(HELIX_DISABLED_PARTITION_STR);
  }
  helixAdmin.setInstanceConfig(clusterName, instanceName, instanceConfig);
  instancesUpdated.getAndIncrement();
  if (helixAdminOperation == HelixAdminOperation.EnablePartition) {
    partitionsEnabled.getAndIncrement();
  } else {
    partitionsDisabled.getAndIncrement();
  }
}
 

/**
 * 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");
}