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

@Test
public void testLegacyEnableDisablePartition() {
  String instanceName = "TestInstanceLegacy";
  String testResourcePrefix = "TestResourceLegacy";
  ZNRecord record = new ZNRecord(instanceName);
  List<String> disabledPartitions = new ArrayList<>(Arrays.asList("1", "2", "3"));
  record.setListField(InstanceConfig.InstanceConfigProperty.HELIX_DISABLED_PARTITION.name(),
      disabledPartitions);
  InstanceConfig instanceConfig = new InstanceConfig(record);
  instanceConfig.setInstanceEnabledForPartition(testResourcePrefix, "2", false);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix).size(), 3);
  Assert.assertEquals(instanceConfig.getRecord()
          .getListField(InstanceConfig.InstanceConfigProperty.HELIX_DISABLED_PARTITION.name()).size(),
      3);
  instanceConfig.setInstanceEnabledForPartition(testResourcePrefix, "2", true);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix).size(), 2);
  Assert.assertEquals(instanceConfig.getRecord()
          .getListField(InstanceConfig.InstanceConfigProperty.HELIX_DISABLED_PARTITION.name()).size(),
      2);
}
 

@Test
public void testEnableDisablePartitions() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String instanceName = "TestInstance";
  String testResourcePrefix = "TestResource";
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addInstance(clusterName, new InstanceConfig(instanceName));

  // Test disable instances with resources
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "0",
      Arrays.asList("1", "2"));
  admin.enablePartition(false, clusterName, instanceName, testResourcePrefix + "1",
      Arrays.asList("2", "3", "4"));
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 2);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 3);

  // Test disable partition across resources
  // TODO : Remove this part once setInstanceEnabledForPartition(partition, enabled) is removed
  instanceConfig.setInstanceEnabledForPartition("10", false);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "0").size(), 3);
  Assert.assertEquals(instanceConfig.getDisabledPartitions(testResourcePrefix + "1").size(), 4);
  admin.dropCluster(clusterName);
}
 

@Test(dependsOnMethods = "testInstancesStoppable_zoneBased")
public void testInstancesStoppable_disableOneInstance() throws IOException {
  // Disable one selected instance0, it should failed to check
  String instance = "instance0";
  InstanceConfig instanceConfig = _configAccessor.getInstanceConfig(STOPPABLE_CLUSTER, instance);
  instanceConfig.setInstanceEnabled(false);
  instanceConfig.setInstanceEnabledForPartition("FakeResource", "FakePartition", false);
  _configAccessor.setInstanceConfig(STOPPABLE_CLUSTER, instance, instanceConfig);

  // It takes time to reflect the changes.
  BestPossibleExternalViewVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(STOPPABLE_CLUSTER).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());

  Entity entity = Entity.entity("", MediaType.APPLICATION_JSON_TYPE);
  Response response = new JerseyUriRequestBuilder("clusters/{}/instances/{}/stoppable")
      .format(STOPPABLE_CLUSTER, instance).post(this, entity);
  JsonNode jsonResult = OBJECT_MAPPER.readTree(response.readEntity(String.class));
  Assert.assertFalse(jsonResult.get("stoppable").asBoolean());
  Assert.assertEquals(getStringSet(jsonResult, "failedChecks"),
          ImmutableSet.of("HELIX:HAS_DISABLED_PARTITION","HELIX:INSTANCE_NOT_ENABLED","HELIX:INSTANCE_NOT_STABLE","HELIX:MIN_ACTIVE_REPLICA_CHECK_FAILED"));

  // Reenable instance0, it should passed the check
  instanceConfig.setInstanceEnabled(true);
  instanceConfig.setInstanceEnabledForPartition("FakeResource", "FakePartition", true);
  _configAccessor.setInstanceConfig(STOPPABLE_CLUSTER, instance, instanceConfig);
  Assert.assertTrue(verifier.verifyByPolling());

  entity = Entity.entity("", MediaType.APPLICATION_JSON_TYPE);
  response = new JerseyUriRequestBuilder("clusters/{}/instances/{}/stoppable")
      .format(STOPPABLE_CLUSTER, instance).post(this, entity);
  jsonResult = OBJECT_MAPPER.readTree(response.readEntity(String.class));

  Assert.assertFalse(jsonResult.get("stoppable").asBoolean());
  Assert.assertEquals(getStringSet(jsonResult, "failedChecks"), ImmutableSet.of("HELIX:MIN_ACTIVE_REPLICA_CHECK_FAILED"));
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

@Override
public void setReplicaDisabledState(ReplicaId replicaId, boolean disable) {
  if (!(replicaId instanceof AmbryReplica)) {
    throw new IllegalArgumentException(
        "HelixParticipant only works with the AmbryReplica implementation of ReplicaId");
  }
  synchronized (helixAdministrationLock) {
    List<String> disabledReplicas = new ArrayList<>(getDisabledReplicas());
    String partitionName = replicaId.getPartitionId().toPathString();

    // 1. invoke Helix native method to enable/disable partition on local node, this will trigger subsequent state
    //    transition on given replica. This method modifies MapFields in InstanceConfig.
    InstanceConfig instanceConfig = getInstanceConfig();
    String resourceNameForPartition = getResourceNameOfPartition(helixAdmin, clusterName, partitionName);
    logger.info("{} replica {} on current node", disable ? "Disabling" : "Enabling", partitionName);
    instanceConfig.setInstanceEnabledForPartition(resourceNameForPartition, partitionName, !disable);

    // 2. update disabled replica list in InstanceConfig. This modifies ListFields only
    if (!disable && disabledReplicas.contains(partitionName)) {
      logger.info("Removing the partition {} from disabledReplicas list", partitionName);
      disabledReplicas.remove(partitionName);
    } else if (disable && !disabledReplicas.contains(partitionName)) {
      logger.info("Adding the partition {} to disabledReplicas list", partitionName);
      disabledReplicas.add(partitionName);
    }
    logger.info("Setting InstanceConfig with list of disabled replicas: {}", disabledReplicas.toArray());
    instanceConfig.getRecord().setListField(DISABLED_REPLICAS_STR, disabledReplicas);

    // 3. set InstanceConfig in Helix to persist replica disabled state
    helixAdmin.setInstanceConfig(clusterName, instanceName, instanceConfig);
    logger.info("Disabled state of partition {} is updated", partitionName);
  }
}
 

protected ResourceControllerDataProvider setupClusterDataCache() throws IOException {
  ResourceControllerDataProvider testCache = Mockito.mock(ResourceControllerDataProvider.class);

  // 1. Set up the default instance information with capacity configuration.
  InstanceConfig testInstanceConfig = createMockInstanceConfig(_testInstanceId);
  testInstanceConfig.setInstanceEnabledForPartition("TestResource", "TestPartition", false);
  Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
  instanceConfigMap.put(_testInstanceId, testInstanceConfig);
  when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);

  // 2. Set up the basic cluster configuration.
  ClusterConfig testClusterConfig = new ClusterConfig("testClusterConfigId");
  testClusterConfig.setMaxPartitionsPerInstance(5);
  testClusterConfig.setDisabledInstances(Collections.emptyMap());
  testClusterConfig.setInstanceCapacityKeys(new ArrayList<>(_capacityDataMap.keySet()));
  testClusterConfig.setDefaultPartitionWeightMap(
      _capacityDataMap.keySet().stream().collect(Collectors.toMap(key -> key, key -> 0)));
  testClusterConfig.setTopologyAwareEnabled(true);
  when(testCache.getClusterConfig()).thenReturn(testClusterConfig);
  when(testCache.getAbnormalStateResolver(any()))
      .thenReturn(MonitoredAbnormalResolver.DUMMY_STATE_RESOLVER);

  // 3. Mock the live instance node for the default instance.
  LiveInstance testLiveInstance = createMockLiveInstance(_testInstanceId);
  Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
  liveInstanceMap.put(_testInstanceId, testLiveInstance);
  when(testCache.getLiveInstances()).thenReturn(liveInstanceMap);

  // 4. Mock two resources, each with 2 partitions on the default instance.
  // The instance will have the following partitions assigned
  // Resource 1:
  // -------------- partition 1 - MASTER
  // -------------- partition 2 - SLAVE
  // Resource 2:
  // -------------- partition 3 - MASTER
  // -------------- partition 4 - SLAVE
  CurrentState testCurrentStateResource1 = Mockito.mock(CurrentState.class);
  Map<String, String> partitionStateMap1 = new HashMap<>();
  partitionStateMap1.put(_partitionNames.get(0), "MASTER");
  partitionStateMap1.put(_partitionNames.get(1), "SLAVE");
  when(testCurrentStateResource1.getResourceName()).thenReturn(_resourceNames.get(0));
  when(testCurrentStateResource1.getPartitionStateMap()).thenReturn(partitionStateMap1);
  when(testCurrentStateResource1.getStateModelDefRef()).thenReturn("MasterSlave");
  when(testCurrentStateResource1.getState(_partitionNames.get(0))).thenReturn("MASTER");
  when(testCurrentStateResource1.getState(_partitionNames.get(1))).thenReturn("SLAVE");
  CurrentState testCurrentStateResource2 = Mockito.mock(CurrentState.class);
  Map<String, String> partitionStateMap2 = new HashMap<>();
  partitionStateMap2.put(_partitionNames.get(2), "MASTER");
  partitionStateMap2.put(_partitionNames.get(3), "SLAVE");
  when(testCurrentStateResource2.getResourceName()).thenReturn(_resourceNames.get(1));
  when(testCurrentStateResource2.getPartitionStateMap()).thenReturn(partitionStateMap2);
  when(testCurrentStateResource2.getStateModelDefRef()).thenReturn("MasterSlave");
  when(testCurrentStateResource2.getState(_partitionNames.get(2))).thenReturn("MASTER");
  when(testCurrentStateResource2.getState(_partitionNames.get(3))).thenReturn("SLAVE");
  Map<String, CurrentState> currentStatemap = new HashMap<>();
  currentStatemap.put(_resourceNames.get(0), testCurrentStateResource1);
  currentStatemap.put(_resourceNames.get(1), testCurrentStateResource2);
  when(testCache.getCurrentState(_testInstanceId, _sessionId)).thenReturn(currentStatemap);

  // 5. Set up the resource config for the two resources with the partition weight.
  Map<String, Integer> capacityDataMapResource1 = new HashMap<>();
  capacityDataMapResource1.put("item1", 3);
  capacityDataMapResource1.put("item2", 6);
  ResourceConfig testResourceConfigResource1 = new ResourceConfig("Resource1");
  testResourceConfigResource1.setPartitionCapacityMap(
      Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY, capacityDataMapResource1));
  when(testCache.getResourceConfig("Resource1")).thenReturn(testResourceConfigResource1);
  Map<String, Integer> capacityDataMapResource2 = new HashMap<>();
  capacityDataMapResource2.put("item1", 5);
  capacityDataMapResource2.put("item2", 10);
  ResourceConfig testResourceConfigResource2 = new ResourceConfig("Resource2");
  testResourceConfigResource2.setPartitionCapacityMap(
      Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY, capacityDataMapResource2));
  when(testCache.getResourceConfig("Resource2")).thenReturn(testResourceConfigResource2);
  Map<String, ResourceConfig> configMap = new HashMap<>();
  configMap.put("Resource1", testResourceConfigResource1);
  configMap.put("Resource2", testResourceConfigResource2);
  when(testCache.getResourceConfigMap()).thenReturn(configMap);

  // 6. Define mock state model
  for (BuiltInStateModelDefinitions bsmd : BuiltInStateModelDefinitions.values()) {
    when(testCache.getStateModelDef(bsmd.name())).thenReturn(bsmd.getStateModelDefinition());
  }

  return testCache;
}
 

/**
 * Control state of partition on certain node. (i.e. DisablePartition, EnablePartition)
 */
private void controlPartitionState() {
  // for now, we only support controlling state of single partition on certain node. Hence, there should be only 1 dc
  // in adminForDc map.
  if (adminForDc.size() != 1) {
    throw new IllegalStateException("The dc count is not 1 for partition state control operation");
  }
  HelixAdmin helixAdmin = adminForDc.values().iterator().next();
  String instanceName;
  if (portNum == null) {
    Optional<DataNodeId> optionalDataNode =
        staticClusterMap.getDataNodeIds().stream().filter(node -> node.getHostname().equals(hostName)).findFirst();
    if (!optionalDataNode.isPresent()) {
      throw new IllegalStateException("Host " + hostName + " is not found in static clustermap");
    }
    DataNode dataNode = (DataNode) optionalDataNode.get();
    instanceName = getInstanceName(dataNode);
  } else {
    instanceName = ClusterMapUtils.getInstanceName(hostName, portNum);
  }
  InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(clusterName, instanceName);
  String resourceNameForPartition = getResourceNameOfPartition(helixAdmin, clusterName, partitionName);
  info("{} partition {} under resource {} on node {}",
      helixAdminOperation == HelixAdminOperation.EnablePartition ? "Enabling" : "Disabling", partitionName,
      resourceNameForPartition, instanceName);
  instanceConfig.setInstanceEnabledForPartition(resourceNameForPartition, partitionName,
      helixAdminOperation == HelixAdminOperation.EnablePartition);
  // clean up the disabled partition entry if it exists and is empty.
  Map<String, String> disabledPartitions =
      instanceConfig.getRecord().getMapFields().get(HELIX_DISABLED_PARTITION_STR);
  if (disabledPartitions != null && disabledPartitions.isEmpty()) {
    instanceConfig.getRecord().getMapFields().remove(HELIX_DISABLED_PARTITION_STR);
  }
  helixAdmin.setInstanceConfig(clusterName, instanceName, instanceConfig);
  instancesUpdated.getAndIncrement();
  if (helixAdminOperation == HelixAdminOperation.EnablePartition) {
    partitionsEnabled.getAndIncrement();
  } else {
    partitionsDisabled.getAndIncrement();
  }
}