Java Code Examples for org.apache.helix.model.InstanceConfig#addTag()

@Override
public void addInstanceTag(String clusterName, String instanceName, String tag) {
  logger
      .info("Add instance tag {} for instance {} in cluster {}.", tag, instanceName, clusterName);
  if (!ZKUtil.isClusterSetup(clusterName, _zkClient)) {
    throw new HelixException("cluster " + clusterName + " is not setup yet");
  }

  if (!ZKUtil.isInstanceSetup(_zkClient, clusterName, instanceName, InstanceType.PARTICIPANT)) {
    throw new HelixException(
        "cluster " + clusterName + " instance " + instanceName + " is not setup yet");
  }
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  InstanceConfig config = accessor.getProperty(keyBuilder.instanceConfig(instanceName));
  config.addTag(tag);
  accessor.setProperty(keyBuilder.instanceConfig(instanceName), config);
}
 

protected void addInstanceConfig(String storageNodeName, int seqNo, int zoneCount, int tagCount) {
  _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  String zone = "zone-" + seqNo % zoneCount;
  String tag = "tag-" + seqNo % tagCount;

  InstanceConfig config =
      _gSetupTool.getClusterManagementTool().getInstanceConfig(CLUSTER_NAME, storageNodeName);
  config.setDomain(
      String.format("DOMAIN=%s,ZONE=%s,INSTANCE=%s", DOMAIN_NAME, zone, storageNodeName));
  config.addTag(tag);
  _gSetupTool.getClusterManagementTool().setInstanceConfig(CLUSTER_NAME, storageNodeName, config);

  _nodeToZoneMap.put(storageNodeName, zone);
  _nodeToTagMap.put(storageNodeName, tag);
  _nodes.add(storageNodeName);
}
 

/**
 * Returns the Helix InstanceConfig for the given instance.
 */
public static InstanceConfig toHelixInstanceConfig(Instance instance) {
  InstanceConfig instanceConfig = InstanceConfig.toInstanceConfig(getHelixInstanceId(instance));
  List<String> tags = instance.getTags();
  if (tags != null) {
    for (String tag : tags) {
      instanceConfig.addTag(tag);
    }
  }
  Map<String, Integer> pools = instance.getPools();
  if (pools != null && !pools.isEmpty()) {
    Map<String, String> mapValue = new TreeMap<>();
    for (Map.Entry<String, Integer> entry : pools.entrySet()) {
      mapValue.put(entry.getKey(), entry.getValue().toString());
    }
    instanceConfig.getRecord().setMapField(POOL_KEY, mapValue);
  }
  return instanceConfig;
}
 

protected InstanceConfig createMockInstanceConfig(String instanceId) {
  InstanceConfig testInstanceConfig = new InstanceConfig(instanceId);
  testInstanceConfig.setInstanceCapacityMap(_capacityDataMap);
  testInstanceConfig.addTag(_testInstanceTags.get(0));
  testInstanceConfig.setInstanceEnabled(true);
  testInstanceConfig.setZoneId(_testFaultZoneId);
  return testInstanceConfig;
}
 

/**
 * Add instance group tag for controller so that pinot controller can be assigned to lead controller resource.
 */
private void addInstanceGroupTagIfNeeded() {
  InstanceConfig instanceConfig = getHelixInstanceConfig(_instanceId);
  // The instanceConfig can be null when connecting as a participant while running from PerfBenchmarkRunner
  if (instanceConfig != null && !instanceConfig.containsTag(Helix.CONTROLLER_INSTANCE)) {
    LOGGER.info("Controller: {} doesn't contain group tag: {}. Adding one.", _instanceId, Helix.CONTROLLER_INSTANCE);
    instanceConfig.addTag(Helix.CONTROLLER_INSTANCE);
    HelixDataAccessor accessor = _helixZkManager.getHelixDataAccessor();
    accessor.setProperty(accessor.keyBuilder().instanceConfig(_instanceId), instanceConfig);
  }
}
 

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

@Test
public void testGetInstances() throws IOException {
  final String clusterName = "TestTagAwareness_testGetResources";
  final String[] TAGS = {
      "tag1", "tag2"
  };
  final String URL_BASE =
      "http://localhost:" + ADMIN_PORT + "/clusters/" + clusterName + "/instances";

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

  // Add 4 participants, each with differint tag characteristics
  InstanceConfig instance1 = new InstanceConfig("localhost_1");
  instance1.addTag(TAGS[0]);
  admin.addInstance(clusterName, instance1);
  InstanceConfig instance2 = new InstanceConfig("localhost_2");
  instance2.addTag(TAGS[1]);
  admin.addInstance(clusterName, instance2);
  InstanceConfig instance3 = new InstanceConfig("localhost_3");
  instance3.addTag(TAGS[0]);
  instance3.addTag(TAGS[1]);
  admin.addInstance(clusterName, instance3);
  InstanceConfig instance4 = new InstanceConfig("localhost_4");
  admin.addInstance(clusterName, instance4);

  // 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);
  ListInstancesWrapper responseWrapper =
      ClusterRepresentationUtil.JsonToObject(ListInstancesWrapper.class,
          response.getEntityAsText());
  Map<String, List<String>> tagInfo = responseWrapper.tagInfo;

  // Ensure tag ownership is reported correctly
  Assert.assertTrue(tagInfo.containsKey(TAGS[0]));
  Assert.assertTrue(tagInfo.containsKey(TAGS[1]));
  Assert.assertTrue(tagInfo.get(TAGS[0]).contains("localhost_1"));
  Assert.assertFalse(tagInfo.get(TAGS[0]).contains("localhost_2"));
  Assert.assertTrue(tagInfo.get(TAGS[0]).contains("localhost_3"));
  Assert.assertFalse(tagInfo.get(TAGS[0]).contains("localhost_4"));
  Assert.assertFalse(tagInfo.get(TAGS[1]).contains("localhost_1"));
  Assert.assertTrue(tagInfo.get(TAGS[1]).contains("localhost_2"));
  Assert.assertTrue(tagInfo.get(TAGS[1]).contains("localhost_3"));
  Assert.assertFalse(tagInfo.get(TAGS[1]).contains("localhost_4"));
}
 

@Test
public void testDefaultOfflineReplicaGroup() {
  int numReplicas = 3;
  TableConfig tableConfig =
      new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setServerTenant(TENANT_NAME)
          .setNumReplicas(numReplicas)
          .setSegmentAssignmentStrategy(AssignmentStrategy.REPLICA_GROUP_SEGMENT_ASSIGNMENT_STRATEGY).build();
  int numInstancesPerPartition = 2;
  tableConfig.getValidationConfig()
      .setReplicaGroupStrategyConfig(new ReplicaGroupStrategyConfig(null, numInstancesPerPartition));
  InstanceAssignmentDriver driver = new InstanceAssignmentDriver(tableConfig);
  int numInstances = 10;
  List<InstanceConfig> instanceConfigs = new ArrayList<>(numInstances);
  for (int i = 0; i < numInstances; i++) {
    InstanceConfig instanceConfig = new InstanceConfig(SERVER_INSTANCE_ID_PREFIX + i);
    instanceConfig.addTag(OFFLINE_TAG);
    instanceConfigs.add(instanceConfig);
  }

  // Instances should be assigned to 3 replica-groups with a round-robin fashion, each with 2 instances
  InstancePartitions instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
  assertEquals(instancePartitions.getNumReplicaGroups(), numReplicas);
  assertEquals(instancePartitions.getNumPartitions(), 1);
  // Instances of index 4 to 7 are not assigned because of the hash-based rotation
  // Math.abs("myTable_OFFLINE".hashCode()) % 10 = 8
  // [i8, i9, i0, i1, i2, i3, i4, i5, i6, i7]
  //  r0  r1  r2  r0  r1  r2
  assertEquals(instancePartitions.getInstances(0, 0),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 8));
  assertEquals(instancePartitions.getInstances(0, 1),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 2, SERVER_INSTANCE_ID_PREFIX + 9));
  assertEquals(instancePartitions.getInstances(0, 2),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));

  String partitionColumnName = "partition";
  int numPartitions = 2;
  tableConfig.getValidationConfig()
      .setReplicaGroupStrategyConfig(new ReplicaGroupStrategyConfig(partitionColumnName, numInstancesPerPartition));
  SegmentPartitionConfig segmentPartitionConfig = new SegmentPartitionConfig(
      Collections.singletonMap(partitionColumnName, new ColumnPartitionConfig("Modulo", numPartitions)));
  tableConfig.getIndexingConfig().setSegmentPartitionConfig(segmentPartitionConfig);

  // Instances should be assigned to 3 replica-groups with a round-robin fashion, each with 3 instances, then these 3
  // instances should be assigned to 2 partitions, each with 2 instances
  instancePartitions = driver.assignInstances(InstancePartitionsType.OFFLINE, instanceConfigs);
  assertEquals(instancePartitions.getNumReplicaGroups(), numReplicas);
  assertEquals(instancePartitions.getNumPartitions(), numPartitions);
  // Instance of index 7 is not assigned because of the hash-based rotation
  // Math.abs("myTable_OFFLINE".hashCode()) % 10 = 8
  // [i8, i9, i0, i1, i2, i3, i4, i5, i6, i7]
  //  r0, r1, r2, r0, r1, r2, r0, r1, r2
  // r0: [i8, i1, i4]
  //      p0, p0, p1
  //      p1
  // r1: [i9, i2, i5]
  //      p0, p0, p1
  //      p1
  // r2: [i0, i3, i6]
  //      p0, p0, p1
  //      p1
  assertEquals(instancePartitions.getInstances(0, 0),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 1, SERVER_INSTANCE_ID_PREFIX + 8));
  assertEquals(instancePartitions.getInstances(1, 0),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 4, SERVER_INSTANCE_ID_PREFIX + 8));
  assertEquals(instancePartitions.getInstances(0, 1),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 2, SERVER_INSTANCE_ID_PREFIX + 9));
  assertEquals(instancePartitions.getInstances(1, 1),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 5, SERVER_INSTANCE_ID_PREFIX + 9));
  assertEquals(instancePartitions.getInstances(0, 2),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 3));
  assertEquals(instancePartitions.getInstances(1, 2),
      Arrays.asList(SERVER_INSTANCE_ID_PREFIX + 0, SERVER_INSTANCE_ID_PREFIX + 6));
}
 

private void updateInstanceConfigIfNeeded(String host, int port) {
  InstanceConfig instanceConfig = _helixAdmin.getInstanceConfig(_helixClusterName, _instanceId);
  boolean needToUpdateInstanceConfig = false;

  // Add default instance tags if not exist
  List<String> instanceTags = instanceConfig.getTags();
  if (instanceTags == null || instanceTags.size() == 0) {
    if (ZKMetadataProvider.getClusterTenantIsolationEnabled(_helixManager.getHelixPropertyStore())) {
      instanceConfig.addTag(TagNameUtils.getOfflineTagForTenant(null));
      instanceConfig.addTag(TagNameUtils.getRealtimeTagForTenant(null));
    } else {
      instanceConfig.addTag(UNTAGGED_SERVER_INSTANCE);
    }
    needToUpdateInstanceConfig = true;
  }

  // Update host and port if needed
  if (!host.equals(instanceConfig.getHostName())) {
    instanceConfig.setHostName(host);
    needToUpdateInstanceConfig = true;
  }
  String portStr = Integer.toString(port);
  if (!portStr.equals(instanceConfig.getPort())) {
    instanceConfig.setPort(portStr);
    needToUpdateInstanceConfig = true;
  }

  if (needToUpdateInstanceConfig) {
    LOGGER.info("Updating instance config for instance: {} with instance tags: {}, host: {}, port: {}", _instanceId,
        instanceTags, host, port);
  } else {
    LOGGER.info("Instance config for instance: {} has instance tags: {}, host: {}, port: {}, no need to update",
        _instanceId, instanceTags, host, port);
    return;
  }

  // NOTE: Use HelixDataAccessor.setProperty() instead of HelixAdmin.setInstanceConfig() because the latter explicitly
  // forbids instance host/port modification
  HelixDataAccessor helixDataAccessor = _helixManager.getHelixDataAccessor();
  Preconditions.checkState(
      helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().instanceConfig(_instanceId), instanceConfig),
      "Failed to update instance config");
}