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

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

/**
 * Test enabledWagedRebalance by checking the rebalancer class name changed.
 */
@Test
public void testEnableWagedRebalance() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Add an IdealState
  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  admin.addResource(clusterName, testResourcePrefix, idealState);

  admin.enableWagedRebalance(clusterName, Collections.singletonList(testResourcePrefix));
  IdealState is = admin.getResourceIdealState(clusterName, testResourcePrefix);
  Assert.assertEquals(is.getRebalancerClassName(), WagedRebalancer.class.getName());
}
 

private static void createBrokerResourceIfNeeded(String helixClusterName, HelixAdmin helixAdmin,
    boolean enableBatchMessageMode) {
  // Add state model definition if needed
  String stateModel =
      PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.PINOT_BROKER_RESOURCE_ONLINE_OFFLINE_STATE_MODEL;
  StateModelDefinition stateModelDef = helixAdmin.getStateModelDef(helixClusterName, stateModel);
  if (stateModelDef == null) {
    LOGGER.info("Adding state model: {}", stateModel);
    helixAdmin.addStateModelDef(helixClusterName, stateModel,
        PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.generatePinotStateModelDefinition());
  }

  // Add broker resource if needed
  if (helixAdmin.getResourceIdealState(helixClusterName, BROKER_RESOURCE_INSTANCE) == null) {
    LOGGER.info("Adding resource: {}", BROKER_RESOURCE_INSTANCE);
    IdealState idealState = new CustomModeISBuilder(BROKER_RESOURCE_INSTANCE).setStateModel(stateModel).build();
    idealState.setBatchMessageMode(enableBatchMessageMode);
    helixAdmin.addResource(helixClusterName, BROKER_RESOURCE_INSTANCE, idealState);
  }
}
 

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

@GET
@Path("{resourceName}/idealState")
public Response getResourceIdealState(@PathParam("clusterId") String clusterId,
    @PathParam("resourceName") String resourceName) {
  HelixAdmin admin = getHelixAdmin();
  IdealState idealState = admin.getResourceIdealState(clusterId, resourceName);
  if (idealState != null) {
    return JSONRepresentation(idealState.getRecord());
  }

  return notFound();
}
 

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

/**
 * Posts new job to cluster
 */
private void scheduleSingleJob(String jobResource, JobConfig jobConfig) {
  HelixAdmin admin = _manager.getClusterManagmentTool();

  IdealState jobIS = admin.getResourceIdealState(_manager.getClusterName(), jobResource);
  if (jobIS != null) {
    LOG.info("Job " + jobResource + " idealstate already exists!");
    return;
  }

  // Set up job resource based on partitions from target resource

  // Create the UserContentStore for the job first
  TaskUtil.createUserContent(_manager.getHelixPropertyStore(), jobResource,
      new ZNRecord(TaskUtil.USER_CONTENT_NODE));

  int numPartitions = jobConfig.getTaskConfigMap().size();
  if (numPartitions == 0) {
    IdealState targetIs =
        admin.getResourceIdealState(_manager.getClusterName(), jobConfig.getTargetResource());
    if (targetIs == null) {
      LOG.warn("Target resource does not exist for job " + jobResource);
      // do not need to fail here, the job will be marked as failure immediately when job starts
      // running.
    } else {
      numPartitions = targetIs.getPartitionSet().size();
    }
  }

  admin.addResource(_manager.getClusterName(), jobResource, numPartitions,
      TaskConstants.STATE_MODEL_NAME);

  HelixDataAccessor accessor = _manager.getHelixDataAccessor();

  // Set the job configuration
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  HelixProperty resourceConfig = new HelixProperty(jobResource);
  resourceConfig.getRecord().getSimpleFields().putAll(jobConfig.getResourceConfigMap());
  Map<String, TaskConfig> taskConfigMap = jobConfig.getTaskConfigMap();
  if (taskConfigMap != null) {
    for (TaskConfig taskConfig : taskConfigMap.values()) {
      resourceConfig.getRecord().setMapField(taskConfig.getId(), taskConfig.getConfigMap());
    }
  }
  accessor.setProperty(keyBuilder.resourceConfig(jobResource), resourceConfig);

  // Push out new ideal state based on number of target partitions
  IdealStateBuilder builder = new CustomModeISBuilder(jobResource);
  builder.setRebalancerMode(IdealState.RebalanceMode.TASK);
  builder.setNumReplica(1);
  builder.setNumPartitions(numPartitions);
  builder.setStateModel(TaskConstants.STATE_MODEL_NAME);

  if (jobConfig.getInstanceGroupTag() != null) {
    builder.setNodeGroup(jobConfig.getInstanceGroupTag());
  }

  if (jobConfig.isDisableExternalView()) {
    builder.disableExternalView();
  }

  jobIS = builder.build();
  for (int i = 0; i < numPartitions; i++) {
    jobIS.getRecord().setListField(jobResource + "_" + i, new ArrayList<>());
    jobIS.getRecord().setMapField(jobResource + "_" + i, new HashMap<>());
  }
  jobIS.setRebalancerClassName(JobRebalancer.class.getName());
  admin.setResourceIdealState(_manager.getClusterName(), jobResource, jobIS);
}
 

/**
 * Ensure that no assignments happen when there are no tagged nodes, but the resource is tagged
 */
@Test
public void testResourceTaggedFirst() throws Exception {
  final int NUM_PARTICIPANTS = 10;
  final int NUM_PARTITIONS = 4;
  final int NUM_REPLICAS = 2;
  final String RESOURCE_NAME = "TestDB0";
  final String TAG = "ASSIGNABLE";

  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
      "MasterSlave", RebalanceMode.FULL_AUTO, // use FULL_AUTO mode to test node tagging
      true); // do rebalance

  // tag the resource
  HelixAdmin helixAdmin = new ZKHelixAdmin(ZK_ADDR);
  IdealState idealState = helixAdmin.getResourceIdealState(clusterName, RESOURCE_NAME);
  idealState.setInstanceGroupTag(TAG);
  helixAdmin.setResourceIdealState(clusterName, RESOURCE_NAME, idealState);

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

  // start participants
  MockParticipantManager[] participants = new MockParticipantManager[NUM_PARTICIPANTS];
  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    final String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  Thread.sleep(1000);
  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new EmptyZkVerifier(clusterName, RESOURCE_NAME));
  Assert.assertTrue(result, "External view and current state must be empty");

  // cleanup
  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    participants[i].syncStop();
  }
  controller.syncStop();
  TestHelper.dropCluster(clusterName, _gZkClient);
}
 

private Map<String, String> computePartitionHealth(String clusterId, String resourceName) {
  HelixAdmin admin = getHelixAdmin();
  IdealState idealState = admin.getResourceIdealState(clusterId, resourceName);
  ExternalView externalView = admin.getResourceExternalView(clusterId, resourceName);
  StateModelDefinition stateModelDef =
      admin.getStateModelDef(clusterId, idealState.getStateModelDefRef());
  String initialState = stateModelDef.getInitialState();
  List<String> statesPriorityList = stateModelDef.getStatesPriorityList();
  statesPriorityList = statesPriorityList.subList(0, statesPriorityList.indexOf(initialState)); // Trim
                                                                                                // stateList
                                                                                                // to
                                                                                                // initialState
                                                                                                // and
                                                                                                // above
  int minActiveReplicas = idealState.getMinActiveReplicas();

  // Start the logic that determines the health status of each partition
  Map<String, String> partitionHealthResult = new HashMap<>();
  Set<String> allPartitionNames = idealState.getPartitionSet();
  if (!allPartitionNames.isEmpty()) {
    for (String partitionName : allPartitionNames) {
      int replicaCount =
          idealState.getReplicaCount(idealState.getPreferenceList(partitionName).size());
      // Simplify expectedStateCountMap by assuming that all instances are available to reduce
      // computation load on this REST endpoint
      LinkedHashMap<String, Integer> expectedStateCountMap =
          stateModelDef.getStateCountMap(replicaCount, replicaCount);
      // Extract all states into Collections from ExternalView
      Map<String, String> stateMapInExternalView = externalView.getStateMap(partitionName);
      Collection<String> allReplicaStatesInExternalView =
          (stateMapInExternalView != null && !stateMapInExternalView.isEmpty())
              ? stateMapInExternalView.values()
              : Collections.<String> emptyList();
      int numActiveReplicasInExternalView = 0;
      HealthStatus status = HealthStatus.HEALTHY;

      // Go through all states that are "active" states (higher priority than InitialState)
      for (int statePriorityIndex = 0; statePriorityIndex < statesPriorityList
          .size(); statePriorityIndex++) {
        String currentState = statesPriorityList.get(statePriorityIndex);
        int currentStateCountInIdealState = expectedStateCountMap.get(currentState);
        int currentStateCountInExternalView =
            Collections.frequency(allReplicaStatesInExternalView, currentState);
        numActiveReplicasInExternalView += currentStateCountInExternalView;
        // Top state counts must match, if not, unhealthy
        if (statePriorityIndex == 0
            && currentStateCountInExternalView != currentStateCountInIdealState) {
          status = HealthStatus.UNHEALTHY;
          break;
        } else if (currentStateCountInExternalView < currentStateCountInIdealState) {
          // For non-top states, if count in ExternalView is less than count in IdealState,
          // partially healthy
          status = HealthStatus.PARTIAL_HEALTHY;
        }
      }
      if (numActiveReplicasInExternalView < minActiveReplicas) {
        // If this partition does not satisfy the number of minimum active replicas, unhealthy
        status = HealthStatus.UNHEALTHY;
      }
      partitionHealthResult.put(partitionName, status.name());
    }
  }
  return partitionHealthResult;
}
 

public static IdealState getBrokerIdealStates(HelixAdmin admin, String clusterName) {
  return admin.getResourceIdealState(clusterName, BROKER_RESOURCE);
}
 

/**
 * Check if the provided list matches the currently persisted preference list
 * @param admin
 * @param clusterName
 * @param resourceName
 * @param partitionName
 * @param expectPreferenceList
 * @return
 */
private boolean preferenceListIsCorrect(HelixAdmin admin, String clusterName, String resourceName,
    String partitionName, List<String> expectPreferenceList) {
  IdealState idealState = admin.getResourceIdealState(clusterName, resourceName);
  List<String> preferenceList = idealState.getPreferenceList(partitionName);
  return expectPreferenceList.equals(preferenceList);
}