Java Code Examples for org.apache.helix.model.IdealState#setInstanceGroupTag()

/**
 * Create an IdealState for a resource that belongs to a resource group We use
 * "resourceGroupName$resourceInstanceTag" as the IdealState znode name to differetiate different
 * resources from the same resourceGroup.
 */
public IdealState createIdealStateForResourceGroup(String resourceGroupName,
    String resourceTag, int numPartition, int replica, String rebalanceMode, String stateModelDefName) {
  String idealStateId = genIdealStateNameWithResourceTag(resourceGroupName, resourceTag);
  IdealState idealState = new IdealState(idealStateId);
  idealState.setNumPartitions(numPartition);
  idealState.setStateModelDefRef(stateModelDefName);
  IdealState.RebalanceMode mode =
      idealState.rebalanceModeFromString(rebalanceMode, IdealState.RebalanceMode.SEMI_AUTO);
  idealState.setRebalanceMode(mode);
  idealState.setReplicas("" + replica);
  idealState.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  idealState.setResourceGroupName(resourceGroupName);
  idealState.setInstanceGroupTag(resourceTag);
  idealState.enableGroupRouting(true);

  return idealState;
}
 

@Test(dependsOnMethods = "test")
public void testWithInstanceTag() throws Exception {
  Set<String> tags = new HashSet<String>(_nodeToTagMap.values());
  int i = 3;
  for (String tag : tags) {
    String db = "Test-DB-" + TestHelper.getTestMethodName() + i++;
    createResourceWithWagedRebalance(CLUSTER_NAME, db,
        BuiltInStateModelDefinitions.MasterSlave.name(), PARTITIONS, _replica, _replica);
    IdealState is =
        _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
    is.setInstanceGroupTag(tag);
    _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, db, is);
    _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db, _replica);
    _allDBs.add(db);
  }
  Thread.sleep(300);
  validate(_replica);
}
 

@Test
public void testZoneIsolationWithInstanceTag() throws Exception {
  Set<String> tags = new HashSet<String>(_nodeToTagMap.values());
  int i = 0;
  for (String tag : tags) {
    String db = "Test-DB-testZoneIsolationWithInstanceTag" + i++;
    createResourceWithWagedRebalance(CLUSTER_NAME, db,
        BuiltInStateModelDefinitions.MasterSlave.name(), PARTITIONS, _replica, _replica);
    IdealState is =
        _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db);
    is.setInstanceGroupTag(tag);
    _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, db, is);
    _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db, _replica);
    _allDBs.add(db);
  }
  Thread.sleep(300);

  validate(_replica);
}
 

@Test(dependsOnMethods = "testParticipantUnavailable")
public void testTagSetIncorrect() throws Exception {
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, testDb, 5,
      BuiltInStateModelDefinitions.MasterSlave.name(), RebalanceMode.FULL_AUTO.name());
  ZkHelixClusterVerifier verifier = new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME)
      .setZkAddr(ZK_ADDR).setResources(new HashSet<>(Collections.singleton(testDb))).build();
  _gSetupTool.getClusterManagementTool().rebalance(CLUSTER_NAME, testDb, 3);
  Assert.assertTrue(verifier.verifyByPolling());

  // Verify there is no rebalance error logged
  Assert.assertNull(accessor.getProperty(errorNodeKey));

  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  checkRebalanceFailureGauge(false);
  checkResourceBestPossibleCalFailureState(ResourceMonitor.RebalanceStatus.NORMAL, testDb);

  // set expected instance tag
  IdealState is =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, testDb);
  is.setInstanceGroupTag("RandomTag");
  _gSetupTool.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME, testDb, is);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, testDb, 3);

  // Verify there is rebalance error logged
  pollForError(accessor, errorNodeKey);
  checkRebalanceFailureGauge(true);
  checkResourceBestPossibleCalFailureState(
      ResourceMonitor.RebalanceStatus.BEST_POSSIBLE_STATE_CAL_FAILED, testDb);

  // clean up
  _gSetupTool.getClusterManagementTool().dropResource(CLUSTER_NAME, testDb);
}
 

protected void setupDBs(ClusterSetup clusterSetup) {
  // Set up target db
  if (_numDbs > 1) {
    for (int i = 0; i < _numDbs; i++) {
      int varyNum = _partitionVary ? 10 * i : 0;
      String db = WorkflowGenerator.DEFAULT_TGT_DB + i;
      clusterSetup.addResourceToCluster(CLUSTER_NAME, db, _numPartitions + varyNum,
          MASTER_SLAVE_STATE_MODEL, IdealState.RebalanceMode.FULL_AUTO.toString());
      clusterSetup.rebalanceStorageCluster(CLUSTER_NAME, db, _numReplicas);
      _testDbs.add(db);
    }
  } else {
    if (_instanceGroupTag) {
      clusterSetup.addResourceToCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
          _numPartitions, "OnlineOffline", IdealState.RebalanceMode.FULL_AUTO.name());
      IdealState idealState = clusterSetup.getClusterManagementTool()
          .getResourceIdealState(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
      idealState.setInstanceGroupTag("TESTTAG0");
      clusterSetup.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME,
          WorkflowGenerator.DEFAULT_TGT_DB, idealState);
    } else {
      clusterSetup.addResourceToCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
          _numPartitions, MASTER_SLAVE_STATE_MODEL, IdealState.RebalanceMode.FULL_AUTO.name());
    }
    clusterSetup.rebalanceStorageCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
        _numReplicas);
  }
}
 

@Test
public void testReadSimpleFieldWithoutDer() throws Exception {
  IdealState idealState = new IdealState("MyDB");
  idealState.setInstanceGroupTag("MyTag");
  String recordStr = ClusterRepresentationUtil.ZNRecordToJson(idealState.getRecord());
  String value =
      ResourceUtil.extractSimpleFieldFromZNRecord(recordStr,
          IdealState.IdealStateProperty.INSTANCE_GROUP_TAG.toString());
  Assert.assertEquals(value, "MyTag");
}
 

@Test
public void testGetResources() throws IOException {
  final String clusterName = "TestTagAwareness_testGetResources";
  final String TAG = "tag";
  final String URL_BASE =
      "http://localhost:" + ADMIN_PORT + "/clusters/" + clusterName + "/resourceGroups";

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

  // Add a tagged resource
  IdealState taggedResource = new IdealState("taggedResource");
  taggedResource.setInstanceGroupTag(TAG);
  taggedResource.setStateModelDefRef("OnlineOffline");
  admin.addResource(clusterName, taggedResource.getId(), taggedResource);

  // Add an untagged resource
  IdealState untaggedResource = new IdealState("untaggedResource");
  untaggedResource.setStateModelDefRef("OnlineOffline");
  admin.addResource(clusterName, untaggedResource.getId(), untaggedResource);

  // 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);
  ZNRecord responseRecord =
      ClusterRepresentationUtil.JsonToObject(ZNRecord.class, response.getEntityAsText());

  // Ensure that the tagged resource has information and the untagged one doesn't
  Assert.assertNotNull(responseRecord.getMapField("ResourceTags"));
  Assert
      .assertEquals(TAG, responseRecord.getMapField("ResourceTags").get(taggedResource.getId()));
  Assert.assertFalse(responseRecord.getMapField("ResourceTags").containsKey(
      untaggedResource.getId()));
}
 

/**
 *
 * Building an empty idealState for a given table.
 * Used when creating a new table.
 *
 * @param tableName resource name
 * @param numCopies is the number of replicas
 * @return
 */
public static IdealState buildEmptyIdealStateFor(String tableName, int numCopies, boolean enableBatchMessageMode) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(tableName);
  final int replicas = numCopies;
  customModeIdealStateBuilder
      .setStateModel(PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL)
      .setNumPartitions(0).setNumReplica(replicas).setMaxPartitionsPerNode(1);
  final IdealState idealState = customModeIdealStateBuilder.build();
  idealState.setInstanceGroupTag(tableName);
  idealState.setBatchMessageMode(enableBatchMessageMode);
  return idealState;
}
 

public static IdealState buildEmptyRealtimeIdealStateFor(String realtimeTableName, int replicaCount,
    boolean enableBatchMessageMode) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(realtimeTableName);
  customModeIdealStateBuilder
      .setStateModel(PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL)
      .setNumPartitions(0).setNumReplica(replicaCount).setMaxPartitionsPerNode(1);
  final IdealState idealState = customModeIdealStateBuilder.build();
  idealState.setInstanceGroupTag(realtimeTableName);
  idealState.setBatchMessageMode(enableBatchMessageMode);

  return idealState;
}
 

void rebalance(String clusterName, String resourceName, int replica, String keyPrefix,
    List<String> instanceNames, String groupId) {
  logger.info("Rebalance resource {} with replica {} in cluster {}.", resourceName, replica,
      clusterName);
  // ensure we get the same idealState with the same set of instances
  Collections.sort(instanceNames);

  IdealState idealState = getResourceIdealState(clusterName, resourceName);
  if (idealState == null) {
    throw new HelixException("Resource: " + resourceName + " has NOT been added yet");
  }

  if (groupId != null && groupId.length() > 0) {
    idealState.setInstanceGroupTag(groupId);
  }
  idealState.setReplicas(Integer.toString(replica));
  int partitions = idealState.getNumPartitions();
  String stateModelName = idealState.getStateModelDefRef();
  StateModelDefinition stateModDef = getStateModelDef(clusterName, stateModelName);

  if (stateModDef == null) {
    throw new HelixException("cannot find state model: " + stateModelName);
  }
  // StateModelDefinition def = new StateModelDefinition(stateModDef);

  List<String> statePriorityList = stateModDef.getStatesPriorityList();

  String masterStateValue = null;
  String slaveStateValue = null;
  replica--;

  for (String state : statePriorityList) {
    String count = stateModDef.getNumInstancesPerState(state);
    if (count.equals("1")) {
      if (masterStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      masterStateValue = state;
    } else if (count.equalsIgnoreCase("R")) {
      if (slaveStateValue != null) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      slaveStateValue = state;
    } else if (count.equalsIgnoreCase("N")) {
      if (!(masterStateValue == null && slaveStateValue == null)) {
        throw new HelixException("Invalid or unsupported state model definition");
      }
      replica = instanceNames.size() - 1;
      masterStateValue = slaveStateValue = state;
    }
  }
  if (masterStateValue == null && slaveStateValue == null) {
    throw new HelixException("Invalid or unsupported state model definition");
  }

  if (masterStateValue == null) {
    masterStateValue = slaveStateValue;
  }
  if (idealState.getRebalanceMode() != RebalanceMode.FULL_AUTO
      && idealState.getRebalanceMode() != RebalanceMode.USER_DEFINED) {
    ZNRecord newIdealState = DefaultIdealStateCalculator
        .calculateIdealState(instanceNames, partitions, replica, keyPrefix, masterStateValue,
            slaveStateValue);

    // for now keep mapField in SEMI_AUTO mode and remove listField in CUSTOMIZED mode
    if (idealState.getRebalanceMode() == RebalanceMode.SEMI_AUTO) {
      idealState.getRecord().setListFields(newIdealState.getListFields());
      // TODO: need consider to remove this.
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
    if (idealState.getRebalanceMode() == RebalanceMode.CUSTOMIZED) {
      idealState.getRecord().setMapFields(newIdealState.getMapFields());
    }
  } else {
    for (int i = 0; i < partitions; i++) {
      String partitionName = keyPrefix + "_" + i;
      idealState.getRecord().setMapField(partitionName, new HashMap<String, String>());
      idealState.getRecord().setListField(partitionName, new ArrayList<String>());
    }
  }
  setResourceIdealState(clusterName, resourceName, idealState);
}
 

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

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