Java Code Examples for org.apache.helix.manager.zk.ZKHelixAdmin#getInstanceConfig()

private static void addInstanceConfig(String instanceName) {
  // add node to cluster if not already added
  ZKHelixAdmin admin = new ZKHelixAdmin(ZK_ADDR);

  InstanceConfig instanceConfig = null;
  try {
    instanceConfig = admin.getInstanceConfig(_clusterName, instanceName);
  } catch (Exception ignored) {
  }
  if (instanceConfig == null) {
    InstanceConfig config = new InstanceConfig(instanceName);
    config.setHostName("localhost");
    config.setInstanceEnabled(true);
    echo("Adding InstanceConfig:" + config);
    admin.addInstance(_clusterName, config);
  }
}
 

/**
 * Initializes the cache loader using the given data source config.
 *
 * @param pinotThirdEyeDataSourceConfig the data source config that provides controller's information.
 *
 * @throws Exception when an error occurs connecting to the Pinot controller.
 */
private void init(PinotThirdEyeDataSourceConfig pinotThirdEyeDataSourceConfig) throws Exception {
  if (pinotThirdEyeDataSourceConfig.getBrokerUrl() != null
      && pinotThirdEyeDataSourceConfig.getBrokerUrl().trim().length() > 0) {
    ZkClient zkClient = new ZkClient(pinotThirdEyeDataSourceConfig.getZookeeperUrl());
    zkClient.setZkSerializer(new ZNRecordSerializer());
    zkClient.waitUntilConnected(CONNECTION_TIMEOUT, TimeUnit.MILLISECONDS);
    ZKHelixAdmin helixAdmin = new ZKHelixAdmin(zkClient);
    List<String> thirdeyeBrokerList = helixAdmin.getInstancesInClusterWithTag(
        pinotThirdEyeDataSourceConfig.getClusterName(), pinotThirdEyeDataSourceConfig.getTag());

    String[] thirdeyeBrokers = new String[thirdeyeBrokerList.size()];
    for (int i = 0; i < thirdeyeBrokerList.size(); i++) {
      String instanceName = thirdeyeBrokerList.get(i);
      InstanceConfig instanceConfig =
          helixAdmin.getInstanceConfig(pinotThirdEyeDataSourceConfig.getClusterName(), instanceName);
      thirdeyeBrokers[i] = instanceConfig.getHostName().replaceAll(BROKER_PREFIX, "") + ":"
          + instanceConfig.getPort();
    }
    this.connections = fromHostList(thirdeyeBrokers);
    LOG.info("Created PinotControllerResponseCacheLoader with brokers {}", thirdeyeBrokers);
  } else {
    this.connections = fromZookeeper(pinotThirdEyeDataSourceConfig);
    LOG.info("Created PinotControllerResponseCacheLoader with controller {}:{}",
        pinotThirdEyeDataSourceConfig.getControllerHost(), pinotThirdEyeDataSourceConfig.getControllerPort());
  }
}
 

/**
 * Test that when AdminOperation is specified to UpdateIdealState, Helix bootstrap tool updates IdealState only without
 * changing InstanceConfig.
 */
@Test
public void testUpdateIdealStateAdminOp() throws Exception {
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  // Test regular bootstrap. This is to ensure InstanceConfig and IdealState are there before testing changing
  // IdealState (to trigger replica movement)
  long expectedResourceCount =
      (testPartitionLayout.getPartitionLayout().getPartitionCount() - 1) / DEFAULT_MAX_PARTITIONS_PER_RESOURCE + 1;
  writeBootstrapOrUpgrade(expectedResourceCount, false);

  // Now, change the replica count for two partitions.
  int totalPartitionCount = testPartitionLayout.getPartitionCount();
  int firstPartitionIndex = RANDOM.nextInt(totalPartitionCount);
  int secondPartitionIndex = (firstPartitionIndex + 1) % totalPartitionCount;
  List<PartitionId> allPartitions = testPartitionLayout.getPartitionLayout().getPartitions(null);
  Partition partition1 = (Partition) allPartitions.get(firstPartitionIndex);
  Partition partition2 = (Partition) allPartitions.get(secondPartitionIndex);

  // Add a new replica for partition1. Find a disk on a data node that does not already have a replica for partition1.
  HashSet<DataNodeId> partition1Nodes = new HashSet<>();
  for (ReplicaId replica : partition1.getReplicas()) {
    partition1Nodes.add(replica.getDataNodeId());
  }
  Disk diskForNewReplica;
  do {
    diskForNewReplica = testHardwareLayout.getRandomDisk();
  } while (partition1Nodes.contains(diskForNewReplica.getDataNode()) || !diskForNewReplica.getDataNode()
      .getDatacenterName()
      .equals("DC1"));
  // Add new replica into partition1
  ReplicaId replicaToAdd = new Replica(partition1, diskForNewReplica, testHardwareLayout.clusterMapConfig);
  partition1.addReplica(replicaToAdd);
  // Remove a replica from partition2.
  ReplicaId removedReplica = partition2.getReplicas().remove(0);

  String dcName = replicaToAdd.getDataNodeId().getDatacenterName();
  ZkInfo zkInfo = dcsToZkInfo.get(dcName);
  ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, getInstanceName(replicaToAdd.getDataNodeId()));
  // deep copy for subsequent verification
  InstanceConfig previousInstanceConfig = new InstanceConfig(instanceConfig.getRecord());
  Utils.writeJsonObjectToFile(zkJson, zkLayoutPath);
  Utils.writeJsonObjectToFile(testHardwareLayout.getHardwareLayout().toJSONObject(), hardwareLayoutPath);
  Utils.writeJsonObjectToFile(testPartitionLayout.getPartitionLayout().toJSONObject(), partitionLayoutPath);
  // upgrade Helix by updating IdealState: AdminOperation = UpdateIdealState
  HelixBootstrapUpgradeUtil.bootstrapOrUpgrade(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
      CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, false, false, new HelixAdminFactory(), false,
      ClusterMapConfig.DEFAULT_STATE_MODEL_DEF, UpdateIdealState);
  verifyResourceCount(testHardwareLayout.getHardwareLayout(), expectedResourceCount);

  // verify IdealState has been updated
  // 1. new added replica is indeed present in the IdealState
  verifyIdealStateForPartition(replicaToAdd, true, 4, expectedResourceCount);
  // 2. removed old replica is no longer present in the IdealState
  verifyIdealStateForPartition(removedReplica, false, 2, expectedResourceCount);

  // verify the InstanceConfig stays unchanged
  InstanceConfig currentInstanceConfig =
      admin.getInstanceConfig(clusterName, getInstanceName(replicaToAdd.getDataNodeId()));
  assertEquals("InstanceConfig should stay unchanged", previousInstanceConfig.getRecord(),
      currentInstanceConfig.getRecord());
}
 

/**
 * Test when AdminOperation is specified to DisablePartition, Helix bootstrap tool is able to disable certain partition
 * only without changing IdealState and InstanceConfig. (In practice, this is first step to decommission a replica)
 * @throws Exception
 */
@Test
public void testDisablePartitionAdminOp() throws Exception {
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  // Test regular bootstrap. This is to ensure InstanceConfig and IdealState are there before testing disabling certain
  // replica on specific node.
  long expectedResourceCount =
      (testPartitionLayout.getPartitionLayout().getPartitionCount() - 1) / DEFAULT_MAX_PARTITIONS_PER_RESOURCE + 1;
  writeBootstrapOrUpgrade(expectedResourceCount, false);
  int totalPartitionCount = testPartitionLayout.getPartitionCount();
  // Randomly pick a partition to remove one of its replicas
  Partition testPartition = (Partition) testPartitionLayout.getPartitionLayout()
      .getPartitions(null)
      .get(RANDOM.nextInt(totalPartitionCount));

  ReplicaId removedReplica = testPartition.getReplicaIds()
      .stream()
      .filter(r -> r.getDataNodeId().getDatacenterName().equals("DC1"))
      .findFirst()
      .get();
  testPartition.getReplicas().remove(removedReplica);
  ZkInfo zkInfo = dcsToZkInfo.get(removedReplica.getDataNodeId().getDatacenterName());
  ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
  InstanceConfig instanceConfig =
      admin.getInstanceConfig(clusterName, getInstanceName(removedReplica.getDataNodeId()));
  // Deep copy the InstanceConfig for validation
  InstanceConfig previousInstanceConfig = new InstanceConfig(instanceConfig.getRecord());
  // Write changes to static files
  Utils.writeJsonObjectToFile(zkJson, zkLayoutPath);
  Utils.writeJsonObjectToFile(testHardwareLayout.getHardwareLayout().toJSONObject(), hardwareLayoutPath);
  Utils.writeJsonObjectToFile(testPartitionLayout.getPartitionLayout().toJSONObject(), partitionLayoutPath);
  // Upgrade Helix by updating IdealState: AdminOperation = DisablePartition
  HelixBootstrapUpgradeUtil.bootstrapOrUpgrade(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
      CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, false, false, new HelixAdminFactory(), false,
      ClusterMapConfig.DEFAULT_STATE_MODEL_DEF, DisablePartition);
  verifyResourceCount(testHardwareLayout.getHardwareLayout(), expectedResourceCount);
  // Verify that IdealState has no change
  verifyIdealStateForPartition(removedReplica, true, 3, expectedResourceCount);
  // Verify the InstanceConfig is changed only in MapFields (Disabled partitions are added to this field)
  InstanceConfig currentInstanceConfig =
      admin.getInstanceConfig(clusterName, getInstanceName(removedReplica.getDataNodeId()));
  String disabledPartitionStr = currentInstanceConfig.getRecord()
      .getMapFields()
      .keySet()
      .stream()
      .filter(k -> !k.startsWith("/mnt"))
      .findFirst()
      .get();
  // Deep copy the current InstanceConfig to remove disabled partitions entry and compare it with previous InstanceConfig
  InstanceConfig currentCopy = new InstanceConfig(currentInstanceConfig.getRecord());
  currentCopy.getRecord().getMapFields().remove(disabledPartitionStr);
  assertEquals("InstanceConfig should stay unchanged after disabled partition entry is removed",
      previousInstanceConfig.getRecord(), currentCopy.getRecord());
  // Verify that replica has been disabled
  String resourceName = null;
  for (String rs : admin.getResourcesInCluster(clusterName)) {
    IdealState is = admin.getResourceIdealState(clusterName, rs);
    if (is.getPartitionSet().contains(removedReplica.getPartitionId().toPathString())) {
      resourceName = rs;
      break;
    }
  }
  List<String> disabledPartition = currentInstanceConfig.getDisabledPartitions(resourceName);
  assertEquals("Disabled partition is not expected",
      Collections.singletonList(removedReplica.getPartitionId().toPathString()), disabledPartition);
}
 

/**
 * Test that partition is correctly enabled on given node. The partition is first disabled and then enabled.
 * @throws Exception
 */
@Test
public void testEnablePartitionAdminOp() throws Exception {
  assumeTrue(!dcStr.equals("DC1") && !dcStr.equals("DC0"));
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  // Test regular bootstrap.
  long expectedResourceCount =
      (testPartitionLayout.getPartitionLayout().getPartitionCount() - 1) / DEFAULT_MAX_PARTITIONS_PER_RESOURCE + 1;
  writeBootstrapOrUpgrade(expectedResourceCount, false);
  int totalPartitionCount = testPartitionLayout.getPartitionCount();
  // Randomly pick a partition to disable/enable
  Partition testPartition = (Partition) testPartitionLayout.getPartitionLayout()
      .getPartitions(null)
      .get(RANDOM.nextInt(totalPartitionCount));
  // Randomly pick a replica from this partition
  List<ReplicaId> replicaIds = testPartition.getReplicaIds();
  DataNodeId dataNodeId = replicaIds.get(RANDOM.nextInt(replicaIds.size())).getDataNodeId();
  // Disable partition on chosen node
  HelixBootstrapUpgradeUtil.controlPartitionState(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
      CLUSTER_NAME_PREFIX, dataNodeId.getDatacenterName(), dataNodeId.getHostname(), dataNodeId.getPort(),
      DisablePartition, testPartition.toPathString());
  // Verify the InstanceConfig is changed only in MapFields (Disabled partition is added to this field)
  ZkInfo zkInfo = dcsToZkInfo.get(dataNodeId.getDatacenterName());
  ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
  InstanceConfig currentInstanceConfig = admin.getInstanceConfig(clusterName, getInstanceName(dataNodeId));
  assertTrue("There should be additional string in InstanceConfig due to disabling partition",
      currentInstanceConfig.getRecord().getMapFields().keySet().stream().anyMatch(k -> !k.startsWith("/mnt")));
  // Verify given partition is indeed disabled on specified node
  String resourceName = getResourceNameOfPartition(admin, clusterName, testPartition.toPathString());
  List<String> disabledPartitions = currentInstanceConfig.getDisabledPartitions(resourceName);
  assertEquals("Disabled partition is not expected", Collections.singletonList(testPartition.toPathString()),
      disabledPartitions);
  // Enable the same partition on same node
  HelixBootstrapUpgradeUtil.controlPartitionState(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
      CLUSTER_NAME_PREFIX, dataNodeId.getDatacenterName(), dataNodeId.getHostname(), dataNodeId.getPort(),
      EnablePartition, testPartition.toPathString());
  // Verify instanceConfig has been updated (disabled partition is removed)
  currentInstanceConfig = admin.getInstanceConfig(clusterName, getInstanceName(dataNodeId));
  assertFalse("There shouldn't be any additional string in InstanceConfig",
      currentInstanceConfig.getRecord().getMapFields().keySet().stream().anyMatch(k -> !k.startsWith("/mnt")));
  // Verify there is no disabled partition
  assertNull("There shouldn't be any disabled partition", currentInstanceConfig.getDisabledPartitions(resourceName));
}
 

/**
 * Test listing sealed partitions in Helix cluster.
 * @throws Exception
 */
@Test
public void testListSealedPartitions() throws Exception {
  assumeTrue(!dcStr.equals("DC1"));
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  // Test regular bootstrap
  long expectedResourceCount =
      (testPartitionLayout.getPartitionLayout().getPartitionCount() - 1) / DEFAULT_MAX_PARTITIONS_PER_RESOURCE + 1;
  writeBootstrapOrUpgrade(expectedResourceCount, false);
  Set<String> sealedPartitions =
      HelixBootstrapUpgradeUtil.listSealedPartition(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
          CLUSTER_NAME_PREFIX, dcStr);
  assertEquals("Sealed partition set should be empty initially", Collections.emptySet(), sealedPartitions);
  // randomly choose 20 partitions to mark as sealed
  int[] intArray = new Random().ints(0, 100).distinct().limit(20).toArray();
  Set<String> selectedSealedPartitionSet = new HashSet<>();
  for (int id : intArray) {
    selectedSealedPartitionSet.add(String.valueOf(id));
  }
  // update the sealed lists in Helix
  for (ZkInfo zkInfo : dcsToZkInfo.values()) {
    ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
    for (String instanceName : admin.getInstancesInCluster(clusterName)) {
      InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
      Set<String> localReplicas = new HashSet<>();
      for (Map<String, String> diskInfo : instanceConfig.getRecord().getMapFields().values()) {
        String replicaStrs = diskInfo.get(REPLICAS_STR);
        for (String replicaStr : replicaStrs.split(REPLICAS_DELIM_STR)) {
          localReplicas.add(replicaStr.split(REPLICAS_STR_SEPARATOR)[0]);
        }
      }
      // derive the intersection of localReplicas set and selectedSealedPartitionSet
      localReplicas.retainAll(selectedSealedPartitionSet);
      instanceConfig.getRecord().setListField(SEALED_STR, new ArrayList<>(localReplicas));
      admin.setInstanceConfig(clusterName, instanceName, instanceConfig);
    }
  }
  // query sealed partition in Helix again
  sealedPartitions =
      HelixBootstrapUpgradeUtil.listSealedPartition(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
          CLUSTER_NAME_PREFIX, dcStr);
  assertEquals("Mismatch in sealed partition set", selectedSealedPartitionSet, sealedPartitions);
}