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

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

private void assignIdealStates(HelixManager helixManager,
    Map<String, IdealState> idealStatesFromAssignment) {
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : idealStatesFromAssignment.keySet()) {
    IdealState idealState = idealStatesFromAssignment.get(topic);
    helixAdmin.setResourceIdealState(helixClusterName, topic, idealState);
  }
}
 

private void startAdmin() {
  HelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  // create cluster
  System.out.println("Creating cluster: " + _clusterName);
  admin.addCluster(_clusterName, true);

  // add MasterSlave state mode definition
  admin.addStateModelDef(_clusterName, "MasterSlave",
      new StateModelDefinition(generateConfigForMasterSlave()));

  // ideal-state znrecord
  ZNRecord record = new ZNRecord(_resourceName);
  record.setSimpleField("IDEAL_STATE_MODE", "AUTO");
  record.setSimpleField("NUM_PARTITIONS", "1");
  record.setSimpleField("REPLICAS", "2");
  record.setSimpleField("STATE_MODEL_DEF_REF", "MasterSlave");
  record.setListField(_resourceName, Arrays.asList("node1", "node2"));

  admin.setResourceIdealState(_clusterName, _resourceName, new IdealState(record));

  ConstraintItemBuilder builder = new ConstraintItemBuilder();

  // limit one transition message at a time across the entire cluster
  builder.addConstraintAttribute("MESSAGE_TYPE", "STATE_TRANSITION")
      // .addConstraintAttribute("INSTANCE", ".*") // un-comment this line if using instance-level
      // constraint
      .addConstraintAttribute("CONSTRAINT_VALUE", "1");
  admin.setConstraint(_clusterName, ClusterConstraints.ConstraintType.MESSAGE_CONSTRAINT,
      "constraint1", builder.build());
}
 

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