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

@Test
public void disabledEmptyTableTest()
    throws Exception {

  final String tableName = "myTable_OFFLINE";
  List<String> allTableNames = Lists.newArrayList(tableName);
  IdealState idealState = new IdealState(tableName);
  // disable table in idealstate
  idealState.enable(false);
  idealState.setReplicas("1");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

  {
    helixResourceManager = mock(PinotHelixResourceManager.class);
    when(helixResourceManager.getAllTables()).thenReturn(allTableNames);
    when(helixResourceManager.getTableIdealState(tableName)).thenReturn(idealState);
    when(helixResourceManager.getTableExternalView(tableName)).thenReturn(null);
  }
  {
    config = mock(ControllerConf.class);
    when(config.getStatusCheckerFrequencyInSeconds()).thenReturn(300);
    when(config.getStatusCheckerWaitForPushTimeInSeconds()).thenReturn(300);
  }
  {
    leadControllerManager = mock(LeadControllerManager.class);
    when(leadControllerManager.isLeaderForTable(anyString())).thenReturn(true);
  }
  metricsRegistry = new MetricsRegistry();
  controllerMetrics = new ControllerMetrics(metricsRegistry);
  segmentStatusChecker =
      new SegmentStatusChecker(helixResourceManager, leadControllerManager, config, controllerMetrics);
  // verify state before test
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 0);
  // update metrics
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 1);
}
 

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

/**
 * Creates a setup where the health API can be tested.
 * @param clusterName
 * @param resourceName
 * @param idealStateParams
 * @param partitionReplicaStates maps partitionName to its replicas' states
 * @throws Exception
 */
private void createDummyMapping(String clusterName, String resourceName,
    Map<String, String> idealStateParams, Map<String, List<String>> partitionReplicaStates)
    throws Exception {
  IdealState idealState = new IdealState(resourceName);
  idealState.setMinActiveReplicas(Integer.parseInt(idealStateParams.get("MinActiveReplicas"))); // 2
  idealState.setStateModelDefRef(idealStateParams.get("StateModelDefRef")); // MasterSlave
  idealState.setMaxPartitionsPerInstance(
      Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance"))); // 3
  idealState.setReplicas(idealStateParams.get("Replicas")); // 3
  idealState.setNumPartitions(Integer.parseInt(idealStateParams.get("NumPartitions"))); // 3
  idealState.enable(false);

  Map<String, List<String>> partitionNames = new LinkedHashMap<>();
  List<String> dummyPrefList = new ArrayList<>();

  for (int i = 0; i < Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance")); i++) {
    dummyPrefList.add(ANY_INSTANCE);
    partitionNames.put("p" + i, dummyPrefList);
  }
  idealState.getRecord().getListFields().putAll(partitionNames);

  if (!_gSetupTool.getClusterManagementTool().getClusters().contains(clusterName)) {
    _gSetupTool.getClusterManagementTool().addCluster(clusterName);
  }
  _gSetupTool.getClusterManagementTool().setResourceIdealState(clusterName, resourceName,
      idealState);

  // Set ExternalView's replica states for a given parameter map
  ExternalView externalView = new ExternalView(resourceName);

  Map<String, Map<String, String>> mappingCurrent = new LinkedHashMap<>();

  List<String> partitionReplicaStatesList = new ArrayList<>(partitionReplicaStates.keySet());
  for (int k = 0; k < partitionReplicaStatesList.size(); k++) {
    Map<String, String> replicaStatesForPartition = new LinkedHashMap<>();
    List<String> replicaStateList = partitionReplicaStates.get(partitionReplicaStatesList.get(k));
    for (int i = 0; i < replicaStateList.size(); i++) {
      replicaStatesForPartition.put("r" + i, replicaStateList.get(i));
    }
    mappingCurrent.put("p" + k, replicaStatesForPartition);
  }

  externalView.getRecord().getMapFields().putAll(mappingCurrent);

  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "p1",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  helixManager.connect();
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().externalView(resourceName),
      externalView);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

@Test
public void disabledTableTest()
    throws Exception {

  final String tableName = "myTable_OFFLINE";
  List<String> allTableNames = new ArrayList<String>();
  allTableNames.add(tableName);
  IdealState idealState = new IdealState(tableName);
  // disable table in idealstate
  idealState.enable(false);
  idealState.setPartitionState("myTable_OFFLINE", "pinot1", "OFFLINE");
  idealState.setPartitionState("myTable_OFFLINE", "pinot2", "OFFLINE");
  idealState.setPartitionState("myTable_OFFLINE", "pinot3", "OFFLINE");
  idealState.setReplicas("1");
  idealState.setRebalanceMode(IdealState.RebalanceMode.CUSTOMIZED);

  {
    helixResourceManager = mock(PinotHelixResourceManager.class);
    when(helixResourceManager.getAllTables()).thenReturn(allTableNames);
    when(helixResourceManager.getTableIdealState(tableName)).thenReturn(idealState);
    when(helixResourceManager.getTableExternalView(tableName)).thenReturn(null);
  }
  {
    config = mock(ControllerConf.class);
    when(config.getStatusCheckerFrequencyInSeconds()).thenReturn(300);
    when(config.getStatusCheckerWaitForPushTimeInSeconds()).thenReturn(300);
  }
  {
    leadControllerManager = mock(LeadControllerManager.class);
    when(leadControllerManager.isLeaderForTable(anyString())).thenReturn(true);
  }
  metricsRegistry = new MetricsRegistry();
  controllerMetrics = new ControllerMetrics(metricsRegistry);
  segmentStatusChecker =
      new SegmentStatusChecker(helixResourceManager, leadControllerManager, config, controllerMetrics);
  // verify state before test
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 0);
  // update metrics
  segmentStatusChecker.start();
  segmentStatusChecker.run();
  Assert.assertEquals(controllerMetrics.getValueOfGlobalGauge(ControllerGauge.DISABLED_TABLE_COUNT), 1);
}