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

@Test
public void testDisableResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");
  String resourceName = "TestDB";
  admin.addStateModelDef(clusterName, "MasterSlave",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave()));
  admin.addResource(clusterName, resourceName, 4, "MasterSlave");
  admin.enableResource(clusterName, resourceName, false);
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertFalse(idealState.isEnabled());
  admin.enableResource(clusterName, resourceName, true);
  idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertTrue(idealState.isEnabled());

  admin.dropCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * Test enabledWagedRebalance by checking the rebalancer class name changed.
 */
@Test
public void testEnableWagedRebalance() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Add an IdealState
  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
  admin.addResource(clusterName, testResourcePrefix, idealState);

  admin.enableWagedRebalance(clusterName, Collections.singletonList(testResourcePrefix));
  IdealState is = admin.getResourceIdealState(clusterName, testResourcePrefix);
  Assert.assertEquals(is.getRebalancerClassName(), WagedRebalancer.class.getName());
}
 

private static void createBrokerResourceIfNeeded(String helixClusterName, HelixAdmin helixAdmin,
    boolean enableBatchMessageMode) {
  // Add state model definition if needed
  String stateModel =
      PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.PINOT_BROKER_RESOURCE_ONLINE_OFFLINE_STATE_MODEL;
  StateModelDefinition stateModelDef = helixAdmin.getStateModelDef(helixClusterName, stateModel);
  if (stateModelDef == null) {
    LOGGER.info("Adding state model: {}", stateModel);
    helixAdmin.addStateModelDef(helixClusterName, stateModel,
        PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.generatePinotStateModelDefinition());
  }

  // Add broker resource if needed
  if (helixAdmin.getResourceIdealState(helixClusterName, BROKER_RESOURCE_INSTANCE) == null) {
    LOGGER.info("Adding resource: {}", BROKER_RESOURCE_INSTANCE);
    IdealState idealState = new CustomModeISBuilder(BROKER_RESOURCE_INSTANCE).setStateModel(stateModel).build();
    idealState.setBatchMessageMode(enableBatchMessageMode);
    helixAdmin.addResource(helixClusterName, BROKER_RESOURCE_INSTANCE, idealState);
  }
}
 

@Test
public void testDropResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  HelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");

  tool.addStateModelDef(clusterName, "MasterSlave",
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave()));
  tool.addResource(clusterName, "test-db", 4, "MasterSlave");
  Map<String, String> resourceConfig = new HashMap<>();
  resourceConfig.put("key1", "value1");
  tool.setConfig(new HelixConfigScopeBuilder(ConfigScopeProperty.RESOURCE).forCluster(clusterName)
      .forResource("test-db").build(), resourceConfig);

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  Assert.assertTrue(_gZkClient.exists(keyBuilder.idealStates("test-db").getPath()),
      "test-db ideal-state should exist");
  Assert.assertTrue(_gZkClient.exists(keyBuilder.resourceConfig("test-db").getPath()),
      "test-db resource config should exist");

  tool.dropResource(clusterName, "test-db");
  Assert.assertFalse(_gZkClient.exists(keyBuilder.idealStates("test-db").getPath()),
      "test-db ideal-state should be dropped");
  Assert.assertFalse(_gZkClient.exists(keyBuilder.resourceConfig("test-db").getPath()),
      "test-db resource config should be dropped");

  tool.dropCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

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

/**
 * Populate info on ZooKeeper server and start {@link HelixControllerManager}.
 * @param zkConnectString zk connect string to zk server.
 * @param vcrClusterName the vcr cluster name.
 * @param clusterMap the {@link ClusterMap} to use.
 * @return the created {@link HelixControllerManager}.
 */
public static HelixControllerManager populateZkInfoAndStartController(String zkConnectString, String vcrClusterName,
    ClusterMap clusterMap) {
  HelixZkClient zkClient = DedicatedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkConnectString), new HelixZkClient.ZkClientConfig());
  try {
    zkClient.setZkSerializer(new ZNRecordSerializer());
    ClusterSetup clusterSetup = new ClusterSetup(zkClient);
    clusterSetup.addCluster(vcrClusterName, true);
    HelixAdmin admin = new ZKHelixAdmin(zkClient);
    // set ALLOW_PARTICIPANT_AUTO_JOIN
    HelixConfigScope configScope = new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.CLUSTER).
        forCluster(vcrClusterName).build();
    Map<String, String> helixClusterProperties = new HashMap<>();
    helixClusterProperties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, String.valueOf(true));
    admin.setConfig(configScope, helixClusterProperties);
    // set PersistBestPossibleAssignment
    ConfigAccessor configAccessor = new ConfigAccessor(zkClient);
    ClusterConfig clusterConfig = configAccessor.getClusterConfig(vcrClusterName);
    clusterConfig.setPersistBestPossibleAssignment(true);
    configAccessor.setClusterConfig(vcrClusterName, clusterConfig);

    FullAutoModeISBuilder builder = new FullAutoModeISBuilder(helixResource);
    builder.setStateModel(LeaderStandbySMD.name);
    for (PartitionId partitionId : clusterMap.getAllPartitionIds(null)) {
      builder.add(partitionId.toPathString());
    }
    builder.setRebalanceStrategy(CrushEdRebalanceStrategy.class.getName());
    IdealState idealState = builder.build();
    admin.addResource(vcrClusterName, helixResource, idealState);
    admin.rebalance(vcrClusterName, helixResource, 3, "", "");
    HelixControllerManager helixControllerManager = new HelixControllerManager(zkConnectString, vcrClusterName);
    helixControllerManager.syncStart();
    return helixControllerManager;
  } finally {
    zkClient.close();
  }
}
 

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

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

/**
 * LockManagerDemo clusterName, numInstances, lockGroupName, numLocks
 * @param args
 * @throws Exception
 */
public static void main(String[] args) throws Exception {
  final String zkAddress = "localhost:2199";
  final String clusterName = "lock-manager-demo";
  final String lockGroupName = "lock-group";
  final int numInstances = 3;
  final int numPartitions = 12;
  final boolean startController = false;
  HelixManager controllerManager = null;
  Thread[] processArray;
  processArray = new Thread[numInstances];
  try {
    startLocalZookeeper(2199);
    HelixAdmin admin = new ZKHelixAdmin(zkAddress);
    admin.addCluster(clusterName, true);
    StateModelConfigGenerator generator = new StateModelConfigGenerator();
    admin.addStateModelDef(clusterName, "OnlineOffline",
        new StateModelDefinition(generator.generateConfigForOnlineOffline()));
    admin.addResource(clusterName, lockGroupName, numPartitions, "OnlineOffline",
        RebalanceMode.FULL_AUTO.toString());
    admin.rebalance(clusterName, lockGroupName, 1);
    for (int i = 0; i < numInstances; i++) {
      final String instanceName = "localhost_" + (12000 + i);
      processArray[i] = new Thread(new Runnable() {

        @Override
        public void run() {
          LockProcess lockProcess = null;

          try {
            lockProcess = new LockProcess(clusterName, zkAddress, instanceName, startController);
            lockProcess.start();
            Thread.currentThread().join();
          } catch (InterruptedException e) {
            System.out.println(instanceName + "Interrupted");
            if (lockProcess != null) {
              lockProcess.stop();
            }
          } catch (Exception e) {
            e.printStackTrace();
          }
        }

      });
      processArray[i].start();
    }
    Thread.sleep(3000);
    controllerManager =
        HelixControllerMain.startHelixController(zkAddress, clusterName, "controller",
            HelixControllerMain.STANDALONE);
    Thread.sleep(5000);
    printStatus(admin, clusterName, lockGroupName);
    System.out.println("Stopping localhost_12000");
    processArray[0].interrupt();
    Thread.sleep(3000);
    printStatus(admin, clusterName, lockGroupName);
    Thread.currentThread().join();
  } catch (Exception e) {
    e.printStackTrace();
  } finally {
    if (controllerManager != null) {
      controllerManager.disconnect();
    }
    for (Thread process : processArray) {
      if (process != null) {
        process.interrupt();
      }
    }
  }
}