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

public static PistachioClusterInfo getClusterInfo() {
    try {
        String zookeeperConnStr = ConfigurationManager.getConfiguration().getString("Pistachio.ZooKeeper.Server");
        ZKHelixAdmin admin = new ZKHelixAdmin(zookeeperConnStr);
        IdealState idealState = admin.getResourceIdealState("PistachiosCluster", "PistachiosResource");
        PistachioClusterInfo info = new PistachioClusterInfo();
        info.numPartitions = idealState.getNumPartitions();
        info.numReplicas = Integer.parseInt(idealState.getReplicas());
        info.hostList = admin.getInstancesInCluster("PistachiosCluster");

        logger.info("num partitions: {}, num Replicas: {}, hostList: {}.", info.numPartitions,
            info.numReplicas, Joiner.on(",").join(info.hostList.toArray()));

        return info;
    } catch (Exception e) {
        logger.info("error getting cluster info", e);
        return null;
    }
}
 

public long getTotalPartition(String resource) {
    if (totalParition == -1) {
        synchronized(totalParition) {
            if (totalParition == -1) {
                ZKHelixAdmin admin = new ZKHelixAdmin(zkAddress);
                IdealState idealState = admin.getResourceIdealState(helixClusterName, resource);
                totalParition = (long)idealState.getNumPartitions();
            }
        }
    }

    return totalParition;
}
 

/***
 * A helper method to verify IdealState for given replica's partition. It checks number of replicas (represented by
 * instances) under certain partition and verifies new added replica exists or removed old replica is no longer present.
 * @param replica the replica to check if it exists (could be either added replica or removed replica)
 * @param shouldExist if {@code true}, it means the given replica is newly added and should exist in IdealState.
 *                    {@code false} otherwise.
 * @param expectedReplicaCountForPartition expected number of replicas under certain partition
 * @param expectedResourceCount expected total number resource count in this cluster.
 */
private void verifyIdealStateForPartition(ReplicaId replica, boolean shouldExist,
    int expectedReplicaCountForPartition, long expectedResourceCount) {
  String dcName = replica.getDataNodeId().getDatacenterName();
  ZkInfo zkInfo = dcsToZkInfo.get(dcName);
  String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
  ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
  if (!activeDcSet.contains(dcName)) {
    Assert.assertFalse("Cluster should not be present, as dc " + dcName + " is not enabled",
        admin.getClusters().contains(clusterName));
  } else {
    List<String> resources = admin.getResourcesInCluster(clusterName);
    assertEquals("Resource count mismatch", expectedResourceCount, resources.size());
    String partitionName = replica.getPartitionId().toPathString();
    boolean resourceFound = false;
    for (String resource : resources) {
      IdealState idealState = admin.getResourceIdealState(clusterName, resource);
      if (idealState.getPartitionSet().contains(partitionName)) {
        resourceFound = true;
        Set<String> instances = idealState.getInstanceSet(partitionName);
        assertEquals("Replica number is not expected", expectedReplicaCountForPartition, instances.size());
        String instanceNameOfNewReplica = getInstanceName(replica.getDataNodeId());
        if (shouldExist) {
          assertTrue("Instance set doesn't include new instance that holds added replica",
              instances.contains(instanceNameOfNewReplica));
        } else {
          assertFalse("Instance that holds deleted replica should not exist in the set",
              instances.contains(instanceNameOfNewReplica));
        }
        break;
      }
    }
    assertTrue("No corresponding resource found for partition " + partitionName, resourceFound);
  }
}
 

public boolean init() {

        boolean initialized = false;

        logger.info("Initializing profile server...........");
        logger.info("do nothing setting {}", doNothing);
        try {
            // open profile store
            Configuration conf = ConfigurationManager.getConfiguration();
        ZKHelixAdmin admin = new ZKHelixAdmin(conf.getString(ZOOKEEPER_SERVER));
        IdealState idealState = admin.getResourceIdealState("PistachiosCluster", "PistachiosResource");
        long totalParition = (long)idealState.getNumPartitions();
            profileStore = new LocalStorageEngine(
                    conf.getString(PROFILE_BASE_DIR),
                    (int)totalParition, 8,
                    conf.getInt("StorageEngine.KC.RecordsPerPartition"),
                    conf.getLong("StorageEngine.KC.MemoryPerPartition"));
            ProcessorRegistry.getInstance().init();
            logger.info("creating helix partition sepctator {} {} {}", conf.getString(ZOOKEEPER_SERVER, "EMPTY"),
                    "PistachiosCluster", conf.getString(PROFILE_HELIX_INSTANCE_ID, "EMPTY"));
            helixPartitionSpectator = HelixPartitionSpectator.getInstance(
                    conf.getString(ZOOKEEPER_SERVER), // zkAddr
                    "PistachiosCluster",
                    NativeUtils.getHostname() //InetAddress.getLocalHost().getHostName() //conf.getString(PROFILE_HELIX_INSTANCE_ID) // instanceName
                    );
            // Partition Manager for line spending
            manager = new HelixPartitionManager<>(
                    conf.getString(ZOOKEEPER_SERVER), // zkAddr
                    "PistachiosCluster",
                    NativeUtils.getHostname() //InetAddress.getLocalHost().getHostName() //conf.getString(PROFILE_HELIX_INSTANCE_ID) // instanceName
                    );
            //manager.start("BootstrapOnlineOffline", new BootstrapOnlineOfflineStateModelFactory(new StorePartitionHandlerFactory()));
            manager.start("MasterSlave", new BootstrapOnlineOfflineStateModelFactory(new StorePartitionHandlerFactory()));
            //            }

            initialized = true;
        } catch (Exception e) {
            logger.error("Failed to initialize ProfileServerModule", e);
        }
        logger.info("Finished initializing profile server...........");

        return initialized;
    }
 

@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  _admin = new ZKHelixAdmin(_gZkClient);
  String namespace = "/" + CLUSTER_NAME;
  if (_gZkClient.exists(namespace)) {
    _gZkClient.deleteRecursively(namespace);
  }

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB1, _PARTITIONS, STATE_MODEL,
      IdealState.RebalanceMode.FULL_AUTO + "");

  IdealState idealState = _admin.getResourceIdealState(CLUSTER_NAME, TEST_DB1);
  idealState.setDisableExternalView(true);
  _admin.setResourceIdealState(CLUSTER_NAME, TEST_DB1, idealState);

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB2, _PARTITIONS, STATE_MODEL,
      IdealState.RebalanceMode.FULL_AUTO + "");

  for (int i = 0; i < NODE_NR; i++) {
    instances[i] = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, instances[i]);
  }

  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB1, _replica);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB2, _replica);

  // start dummy participants
  for (int i = 0; i < NODE_NR; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    MockParticipantManager participant =
        new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    participant.syncStart();
    _participants[i] = participant;

  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  boolean result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);
}
 

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