Java Code Examples for org.apache.helix.HelixDataAccessor#getChildValues()

private List<HelixProperty> getProperty(HelixDataAccessor accessor, String dataSpec,
    PropertyKey propertyKeys, PropertyKey propertyKey, String dataType) {
  List<HelixProperty> properties;
  if (Strings.isNullOrEmpty(dataSpec) || dataSpec.equals(MATCH_ALL_SYM)) {
    // TODO: Apply strict check on the getChildValues() call.
    // TODO: For backward compatibility, allow partial read for now. This may reduce the
    // TODO: matches eventually.
    properties = accessor.getChildValues(propertyKeys, false);
  } else {
    HelixProperty data = accessor.getProperty(propertyKey);
    if (data == null) {
      throw new HelixException(
          String.format("Specified %s %s is not found!", dataType, dataSpec));
    }
    properties = Collections.singletonList(data);
  }
  return properties;
}
 

@Override
@PreFetch(enabled = false)
public void onExternalViewChange(List<ExternalView> externalViewList,
    NotificationContext changeContext) {
  HelixConstants.ChangeType changeType = changeContext.getChangeType();
  if (changeType != null && !_sourceDataTypeMap.containsKey(changeType.getPropertyType())) {
    logger.warn(
        "onExternalViewChange called with mismatched change types. Source data types does not contain changed data type: {}",
        changeType);
    return;
  }
  // Refresh with full list of external view.
  if (externalViewList != null && externalViewList.size() > 0) {
    // keep this here for back-compatibility, application can call onExternalViewChange directly
    // with externalview list supplied.
    String keyReference = generateReferenceKey(PropertyType.EXTERNALVIEW.name(),  DEFAULT_STATE_TYPE);
    HelixDataAccessor accessor = changeContext.getManager().getHelixDataAccessor();
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    List<InstanceConfig> configList = accessor.getChildValues(keyBuilder.instanceConfigs(), true);
    List<LiveInstance> liveInstances = accessor.getChildValues(keyBuilder.liveInstances(), true);
    refreshExternalView(externalViewList, configList, liveInstances, keyReference);
  } else {
    ClusterEventType eventType;
    if (_sourceDataTypeMap.containsKey(PropertyType.EXTERNALVIEW)) {
      eventType = ClusterEventType.ExternalViewChange;
    } else if (_sourceDataTypeMap.containsKey(PropertyType.TARGETEXTERNALVIEW)) {
      eventType = ClusterEventType.TargetExternalViewChange;
    } else {
      logger.warn(
          "onExternalViewChange called with mismatched change types. Source data types does not contain changed data type: {}",
          changeType);
      return;
    }
    _routerUpdater.queueEvent(changeContext, eventType, changeType);
  }
}
 

@Override
public void resetInstance(String clusterName, List<String> instanceNames) {
  // TODO: not mp-safe
  logger.info("Reset instances {} in cluster {}.",
      instanceNames == null ? "NULL" : HelixUtil.serializeByComma(instanceNames), clusterName);
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  List<ExternalView> extViews = accessor.getChildValues(keyBuilder.externalViews(), true);

  Set<String> resetInstanceNames = new HashSet<String>(instanceNames);
  for (String instanceName : resetInstanceNames) {
    List<String> resetPartitionNames = new ArrayList<String>();
    for (ExternalView extView : extViews) {
      Map<String, Map<String, String>> stateMap = extView.getRecord().getMapFields();
      for (String partitionName : stateMap.keySet()) {
        Map<String, String> instanceStateMap = stateMap.get(partitionName);

        if (instanceStateMap.containsKey(instanceName) && instanceStateMap.get(instanceName)
            .equals(HelixDefinedState.ERROR.toString())) {
          resetPartitionNames.add(partitionName);
        }
      }
      resetPartition(clusterName, instanceName, extView.getResourceName(), resetPartitionNames);
    }
  }
}
 

@Override
public boolean enableWagedRebalance(String clusterName, List<String> resourceNames) {
  // Null checks
  if (clusterName == null || clusterName.isEmpty()) {
    throw new HelixException("Cluster name is invalid!");
  }
  if (resourceNames == null || resourceNames.isEmpty()) {
    throw new HelixException("Resource name list is invalid!");
  }

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  List<IdealState> idealStates = accessor.getChildValues(keyBuilder.idealStates(), true);
  List<String> nullIdealStates = new ArrayList<>();
  for (int i = 0; i < idealStates.size(); i++) {
    if (idealStates.get(i) == null) {
      nullIdealStates.add(resourceNames.get(i));
    } else {
      idealStates.get(i).setRebalancerClassName(WagedRebalancer.class.getName());
      idealStates.get(i).setRebalanceMode(RebalanceMode.FULL_AUTO);
    }
  }
  if (!nullIdealStates.isEmpty()) {
    throw new HelixException(
        String.format("Not all IdealStates exist in the cluster: %s", nullIdealStates));
  }
  List<PropertyKey> idealStateKeys = new ArrayList<>();
  idealStates.forEach(
      idealState -> idealStateKeys.add(keyBuilder.idealStates(idealState.getResourceName())));
  boolean[] success = accessor.setChildren(idealStateKeys, idealStates);
  for (boolean s : success) {
    if (!s) {
      return false;
    }
  }
  return true;
}
 

@Override
public void run() {
  try {
    if (_shouldRefreshCacheOption.orElse(
        _clusterEventType.equals(ClusterEventType.PeriodicalRebalance) || _clusterEventType
            .equals(ClusterEventType.OnDemandRebalance))) {
      requestDataProvidersFullRefresh();

      HelixDataAccessor accessor = _manager.getHelixDataAccessor();
      PropertyKey.Builder keyBuilder = accessor.keyBuilder();
      List<LiveInstance> liveInstances =
          accessor.getChildValues(keyBuilder.liveInstances(), true);

      if (liveInstances != null && !liveInstances.isEmpty()) {
        NotificationContext changeContext = new NotificationContext(_manager);
        changeContext.setType(NotificationContext.Type.CALLBACK);
        synchronized (_manager) {
          checkLiveInstancesObservation(liveInstances, changeContext);
        }
      }
    }
    forceRebalance(_manager, _clusterEventType);
  } catch (Throwable ex) {
    logger.error("Time task failed. Rebalance task type: " + _clusterEventType + ", cluster: "
        + _clusterName, ex);
  }
}
 

private List<Message> generateMessagesForParticipant(Criteria recipientCriteria, Message message,
    HelixDataAccessor targetDataAccessor) {
  List<Message> messages = new ArrayList<Message>();
  List<Map<String, String>> matchedList =
      _evaluator.evaluateCriteria(recipientCriteria, targetDataAccessor);

  if (!matchedList.isEmpty()) {
    Map<String, String> sessionIdMap = new HashMap<String, String>();
    if (recipientCriteria.isSessionSpecific()) {
      Builder keyBuilder = targetDataAccessor.keyBuilder();
      // For backward compatibility, allow partial read for the live instances.
      // Note that this may cause the pending message to be sent with null target session Id.
      List<LiveInstance> liveInstances =
          targetDataAccessor.getChildValues(keyBuilder.liveInstances(), false);

      for (LiveInstance liveInstance : liveInstances) {
        sessionIdMap.put(liveInstance.getInstanceName(), liveInstance.getEphemeralOwner());
      }
    }
    for (Map<String, String> map : matchedList) {
      String id = UUID.randomUUID().toString();
      Message newMessage = new Message(message.getRecord(), id);
      String srcInstanceName = _manager.getInstanceName();
      String tgtInstanceName = map.get("instanceName");
      // Don't send message to self
      if (recipientCriteria.isSelfExcluded() && srcInstanceName.equalsIgnoreCase(tgtInstanceName)) {
        continue;
      }
      newMessage.setSrcName(srcInstanceName);
      newMessage.setTgtName(tgtInstanceName);
      newMessage.setResourceName(map.get("resourceName"));
      newMessage.setPartitionName(map.get("partitionName"));
      if (recipientCriteria.isSessionSpecific()) {
        newMessage.setTgtSessionId(sessionIdMap.get(tgtInstanceName));
      }
      messages.add(newMessage);
    }
  }
  return messages;
}
 

private void refreshCache() {
  Builder propertyKeyBuilder = new PropertyKey.Builder(cluster);
  HelixDataAccessor helixDataAccessor = admin.getHelixDataAccessor();
  List<LiveInstance> liveInstances =
      helixDataAccessor.getChildValues(propertyKeyBuilder.liveInstances());
  refreshCache(liveInstances);
}
 

/**
 * Use HelixUtil.getIdealAssignmentForWagedFullAuto() to compute the cluster-wide assignment and
 * verify that it matches with the result from the original WAGED rebalancer's algorithm result.
 */
@Test(dependsOnMethods = "test")
public void testRebalanceTool() throws InterruptedException {
  // Create resources for testing
  int i = 0;
  for (String stateModel : _testModels) {
    String db = "Test-DB-" + TestHelper.getTestMethodName() + i++;
    createResourceWithWagedRebalance(CLUSTER_NAME, db, stateModel, PARTITIONS, _replica,
        _replica);
    _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db, _replica);
    _allDBs.add(db);
  }
  Thread.sleep(300);

  validate(_replica);

  // Read cluster parameters from ZK
  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  ClusterConfig clusterConfig =
      dataAccessor.getProperty(dataAccessor.keyBuilder().clusterConfig());
  List<InstanceConfig> instanceConfigs =
      dataAccessor.getChildValues(dataAccessor.keyBuilder().instanceConfigs(), true);
  List<String> liveInstances =
      dataAccessor.getChildNames(dataAccessor.keyBuilder().liveInstances());
  List<IdealState> idealStates =
      dataAccessor.getChildValues(dataAccessor.keyBuilder().idealStates(), true);
  List<ResourceConfig> resourceConfigs =
      dataAccessor.getChildValues(dataAccessor.keyBuilder().resourceConfigs(), true);

  // Verify that utilResult contains the assignment for the resources added
  Map<String, ResourceAssignment> utilResult = HelixUtil
      .getIdealAssignmentForWagedFullAuto(ZK_ADDR, clusterConfig, instanceConfigs, liveInstances,
          idealStates, resourceConfigs);
  Assert.assertNotNull(utilResult);
  Assert.assertEquals(utilResult.size(), _allDBs.size());
  for (IdealState idealState : idealStates) {
    Assert.assertTrue(utilResult.containsKey(idealState.getResourceName()));
    Assert.assertEquals(utilResult.get(idealState.getResourceName()).getRecord().getMapFields(),
        idealState.getRecord().getMapFields());
  }

  // Try to add a few extra instances
  String instance_0 = "instance_0";
  String instance_1 = "instance_1";
  Set<String> newInstances = new HashSet<>();
  newInstances.add(instance_0);
  newInstances.add(instance_1);
  liveInstances.addAll(newInstances);
  for (String instance : newInstances) {
    InstanceConfig instanceConfig = new InstanceConfig(instance);
    instanceConfigs.add(instanceConfig);
  }

  utilResult = HelixUtil
      .getIdealAssignmentForWagedFullAuto(ZK_ADDR, clusterConfig, instanceConfigs, liveInstances,
          idealStates, resourceConfigs);

  Set<String> instancesWithAssignments = new HashSet<>();
  utilResult.values().forEach(
      resourceAssignment -> resourceAssignment.getRecord().getMapFields().values()
          .forEach(entry -> instancesWithAssignments.addAll(entry.keySet())));
  // The newly added instances should contain some partitions
  Assert.assertTrue(instancesWithAssignments.contains(instance_0));
  Assert.assertTrue(instancesWithAssignments.contains(instance_1));
}
 

@Test
public void testCachedCustomizedViews() throws Exception {
  String clusterName = "CLUSTER_" + TestHelper.getTestMethodName();

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  HelixManager manager = new DummyClusterManager(clusterName, accessor);

  // ideal state: node0 is MASTER, node1 is SLAVE
  // replica=2 means 1 master and 1 slave
  setupIdealState(clusterName, new int[] {
      0, 1
  }, new String[] {
      "TestDB"
  }, 1, 2);
  setupLiveInstances(clusterName, new int[] {
      0, 1
  });
  setupStateModel(clusterName);

  ClusterEvent event = new ClusterEvent(ClusterEventType.Unknown);
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider(clusterName);
  event.addAttribute(AttributeName.helixmanager.name(), manager);
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);

  CustomizedStateConfig config = new CustomizedStateConfig();
  List<String> aggregationEnabledTypes = new ArrayList<>();
  aggregationEnabledTypes.add(CUSTOMIZED_STATE_NAME);
  config.setAggregationEnabledTypes(aggregationEnabledTypes);

  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.customizedStateConfig(), config);

  CustomizedState customizedState = new CustomizedState(RESOURCE_NAME);
  customizedState.setState(PARTITION_NAME, "STARTED");
  accessor
      .setProperty(keyBuilder.customizedState(INSTANCE_NAME, "customizedState1", RESOURCE_NAME),
          customizedState);

  CustomizedViewAggregationStage customizedViewComputeStage =
      new CustomizedViewAggregationStage();
  Pipeline dataRefresh = new Pipeline();
  dataRefresh.addStage(new ReadClusterDataStage());
  runPipeline(event, dataRefresh);
  runStage(event, new ResourceComputationStage());
  runStage(event, new CustomizedStateComputationStage());
  runStage(event, customizedViewComputeStage);
  Assert.assertEquals(cache.getCustomizedViewCacheMap().size(),
      accessor.getChildNames(accessor.keyBuilder().customizedViews()).size());

  // Assure there is no customized view got updated when running the stage again
  List<CustomizedView> oldCustomizedViews =
      accessor.getChildValues(accessor.keyBuilder().customizedViews(), true);
  runStage(event, customizedViewComputeStage);
  List<CustomizedView> newCustomizedViews =
      accessor.getChildValues(accessor.keyBuilder().customizedViews(), true);
  Assert.assertEquals(oldCustomizedViews, newCustomizedViews);
  for (int i = 0; i < oldCustomizedViews.size(); i++) {
    Assert.assertEquals(oldCustomizedViews.get(i).getStat().getVersion(),
        newCustomizedViews.get(i).getStat().getVersion());
  }

  if (manager.isConnected()) {
    manager.disconnect(); // For DummyClusterManager, this is not necessary
  }
  deleteLiveInstances(clusterName);
  deleteCluster(clusterName);
}
 

@Test
public void testCachedExternalViews() throws Exception {
  String clusterName = "CLUSTER_" + TestHelper.getTestMethodName();

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  HelixManager manager = new DummyClusterManager(clusterName, accessor);

  // ideal state: node0 is MASTER, node1 is SLAVE
  // replica=2 means 1 master and 1 slave
  setupIdealState(clusterName, new int[] {
      0, 1
  }, new String[] {
      "TestDB"
  }, 1, 2);
  setupLiveInstances(clusterName, new int[] {
      0, 1
  });
  setupStateModel(clusterName);

  ClusterEvent event = new ClusterEvent(ClusterEventType.Unknown);
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider(clusterName);
  event.addAttribute(AttributeName.helixmanager.name(), manager);
  event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);

  ExternalViewComputeStage externalViewComputeStage = new ExternalViewComputeStage();
  Pipeline dataRefresh = new Pipeline();
  dataRefresh.addStage(new ReadClusterDataStage());
  runPipeline(event, dataRefresh);
  runStage(event, new ResourceComputationStage());
  runStage(event, new CurrentStateComputationStage());
  runStage(event, externalViewComputeStage);
  Assert.assertEquals(cache.getExternalViews().values(),
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true));

  // Assure there is no external got updated
  List<ExternalView> oldExternalViews =
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true);
  runStage(event, externalViewComputeStage);
  List<ExternalView> newExternalViews =
      accessor.getChildValues(accessor.keyBuilder().externalViews(), true);
  Assert.assertEquals(oldExternalViews, newExternalViews);
  for (int i = 0; i < oldExternalViews.size(); i++) {
    Assert.assertEquals(oldExternalViews.get(i).getStat().getVersion(),
        newExternalViews.get(i).getStat().getVersion());
  }

  if (manager.isConnected()) {
    manager.disconnect(); // For DummyClusterManager, this is not necessary
  }
  deleteLiveInstances(clusterName);
  deleteCluster(clusterName);
}
 

/**
 * Returns the config for all the instances in the cluster.
 */
public static List<InstanceConfig> getInstanceConfigs(HelixManager helixManager) {
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  return helixDataAccessor.getChildValues(helixDataAccessor.keyBuilder().instanceConfigs());
}