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

public static Set<TopicPartition> getUnassignedPartitions(HelixManager helixManager) {
  Set<TopicPartition> unassignedPartitions = new HashSet<TopicPartition>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    int numPartitions = is.getNumPartitions();
    for (int partition = 0; partition < numPartitions; ++partition) {
      if (is.getInstanceSet(Integer.toString(partition)).isEmpty()) {
        TopicPartition tpi = new TopicPartition(topic, partition);
        unassignedPartitions.add(tpi);
      }
    }
  }
  return unassignedPartitions;
}
 

/**
 * Update the resources in the destination cluster with the new IdealState settings.
 * @param destZkString the destination Zookeeper server string.
 * @param destClusterName the destination cluster name.
 * @param dryRun run without actual change.
 */
static void updateResourceIdealState(String destZkString, String destClusterName, boolean dryRun) {
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : destResources) {
    IdealState currentIdealState = destAdmin.getResourceIdealState(destClusterName, resource);
    IdealState newIdealState = buildIdealState(resource, currentIdealState.getPartitionSet());
    if (dryRun) {
      System.out.println("Will update " + resource + " to new ideal state " + newIdealState.toString());
    } else {
      destAdmin.setResourceIdealState(destClusterName, resource, newIdealState);
      System.out.println("Updated the ideal state for resource " + resource);
      destAdmin.rebalance(destClusterName, resource, REPLICA_NUMBER, "", "");
      System.out.println("Rebalanced resource " + resource + " with REPLICA_NUM: " + REPLICA_NUMBER);
    }
  }
}
 

/**
 * From IdealStates.
 *
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager,
    Map<String, KafkaBrokerTopicObserver> clusterToObserverMap) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap = new HashMap<>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    for (String partition : is.getPartitionSet()) {
      TopicPartition tpi;
      if (partition.startsWith("@")) {
        if (clusterToObserverMap != null) {
          TopicPartition topicParition = clusterToObserverMap.get(getSrcFromRoute(partition))
              .getTopicPartitionWithRefresh(topic);
          int trueNumPartition = topicParition != null ? topicParition.getPartition() : -1;
          tpi = new TopicPartition(topic, trueNumPartition, partition);
        } else {
          tpi = new TopicPartition(topic, -1, partition);
        }
      } else {
        // route
        tpi = new TopicPartition(topic, Integer.parseInt(partition));
      }
      for (String instance : is.getInstanceSet(partition)) {
        instanceToNumTopicPartitionMap.putIfAbsent(instance, new HashSet<>());
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

private List<ExternalView> getExternalViews() {
  String clusterName = helixManager.getClusterName();
  List<ExternalView> externalViewList = new ArrayList<ExternalView>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  List<String> resourcesInCluster = helixAdmin.getResourcesInCluster(clusterName);
  for (String resourceName : resourcesInCluster) {
    ExternalView ev =
        helixManager.getClusterManagmentTool().getResourceExternalView(clusterName, resourceName);
    if (ev != null) {
      externalViewList.add(ev);
    }
  }

  return externalViewList;
}
 

/**
 * Get resource name associated with given partition.
 * @param helixAdmin the {@link HelixAdmin} to access resources in cluster
 * @param clusterName the name of cluster in which the partition resides
 * @param partitionName name of partition
 * @return resource name associated with given partition. {@code null} if not found.
 */
static String getResourceNameOfPartition(HelixAdmin helixAdmin, String clusterName, String partitionName) {
  String result = null;
  for (String resourceName : helixAdmin.getResourcesInCluster(clusterName)) {
    IdealState idealState = helixAdmin.getResourceIdealState(clusterName, resourceName);
    if (idealState.getPartitionSet().contains(partitionName)) {
      result = resourceName;
      break;
    }
  }
  return result;
}
 

/**
 * A method to verify resources and partitions in src cluster and dest cluster are same.
 */
private boolean isSrcDestSync(String srcZkString, String srcClusterName, String destZkString,
    String destClusterName) {

  HelixAdmin srcAdmin = new ZKHelixAdmin(srcZkString);
  Set<String> srcResources = new HashSet<>(srcAdmin.getResourcesInCluster(srcClusterName));
  HelixAdmin destAdmin = new ZKHelixAdmin(destZkString);
  Set<String> destResources = new HashSet<>(destAdmin.getResourcesInCluster(destClusterName));

  for (String resource : srcResources) {
    if (HelixVcrPopulateTool.ignoreResourceKeyWords.stream().anyMatch(resource::contains)) {
      System.out.println("Resource " + resource + " from src cluster is ignored");
      continue;
    }
    if (destResources.contains(resource)) {
      // check if every partition exist.
      Set<String> srcPartitions = srcAdmin.getResourceIdealState(srcClusterName, resource).getPartitionSet();
      Set<String> destPartitions = destAdmin.getResourceIdealState(destClusterName, resource).getPartitionSet();
      for (String partition : srcPartitions) {
        if (!destPartitions.contains(partition)) {
          return false;
        }
      }
    } else {
      return false;
    }
  }
  return true;
}
 

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