org.apache.helix.model.ConstraintItem Java Examples

private static void setConstraints(ZKHelixAdmin admin, int numPartitions) {

        logger.info("pause cluster");
        admin.enableCluster("PistachiosCluster", false);
        // setting partition constraints
        logger.info("setting per partition state transition constraints to 1");
        try {
            for (int constraintId = 0; constraintId < numPartitions; constraintId++) {
                java.util.HashMap<ConstraintAttribute, String>  attributes = new java.util.HashMap<ConstraintAttribute, String>();
                attributes.put(ConstraintAttribute.RESOURCE, "PistachiosResource");
                attributes.put(ConstraintAttribute.PARTITION, "PistachiosResource_"+constraintId);
                logger.info("setting per partition for {} state transition constraints to 1", "PistachiosResource_"+constraintId);
                admin.setConstraint("PistachiosCluster", ConstraintType.STATE_CONSTRAINT, "PistachiosPartitionTransitionConstraint" + constraintId,
                   new ConstraintItem(attributes,"1"));
            }

        } catch(Exception e) {
            logger.info("setting state transition constraints error, roll back and exit", e);
        }
        logger.info("resume cluster");
        admin.enableCluster("PistachiosCluster", true);
  }
 

@Override
public void setConstraint(String clusterName, final ConstraintType constraintType,
    final String constraintId, final ConstraintItem constraintItem) {
  logger.info("Set constraint type {} with constraint id {} for cluster {}.", constraintType,
      constraintId, clusterName);
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_zkClient);

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  String path = keyBuilder.constraint(constraintType.toString()).getPath();

  baseAccessor.update(path, new DataUpdater<ZNRecord>() {
    @Override
    public ZNRecord update(ZNRecord currentData) {
      ClusterConstraints constraints =
          currentData == null ? new ClusterConstraints(constraintType)
              : new ClusterConstraints(currentData);

      constraints.addConstraintItem(constraintId, constraintItem);
      return constraints.getRecord();
    }
  }, AccessOption.PERSISTENT);
}
 

/**
 * set constraint
 * @param clusterName
 * @param constraintType
 * @param constraintId
 * @param constraintAttributesMap : csv-formated constraint key-value pairs
 */
public void setConstraint(String clusterName, String constraintType, String constraintId,
    String constraintAttributesMap) {
  if (clusterName == null || constraintType == null || constraintId == null
      || constraintAttributesMap == null) {
    throw new IllegalArgumentException(
        "fail to set constraint. missing clusterName|constraintType|constraintId|constraintAttributesMap");
  }

  ConstraintType type = ConstraintType.valueOf(constraintType);
  ConstraintItemBuilder builder = new ConstraintItemBuilder();
  Map<String, String> constraintAttributes =
      HelixUtil.parseCsvFormatedKeyValuePairs(constraintAttributesMap);
  ConstraintItem constraintItem = builder.addConstraintAttributes(constraintAttributes).build();
  _admin.setConstraint(clusterName, type, constraintId, constraintItem);
}
 

/**
 * constraints are selected in the order of the following rules: 1) don't select
 * constraints with CONSTRAINT_VALUE=ANY; 2) if one constraint is more specific than the
 * other, select the most specific one 3) if a message matches multiple constraints of
 * incomparable specificity, select the one with the minimum value 4) if a message
 * matches multiple constraints of incomparable specificity, and they all have the same
 * value, select the first in alphabetic order
 */
Set<ConstraintItem> selectConstraints(Set<ConstraintItem> items,
    Map<ConstraintAttribute, String> attributes) {
  Map<String, ConstraintItem> selectedItems = new HashMap<String, ConstraintItem>();
  for (ConstraintItem item : items) {
    // don't select constraints with CONSTRAINT_VALUE=ANY
    if (item.getConstraintValue().equals(ConstraintValue.ANY.toString())) {
      continue;
    }

    String key = item.filter(attributes).toString();
    if (!selectedItems.containsKey(key)) {
      selectedItems.put(key, item);
    } else {
      ConstraintItem existingItem = selectedItems.get(key);
      if (existingItem.match(item.getAttributes())) {
        // item is more specific than existingItem
        selectedItems.put(key, item);
      } else if (!item.match(existingItem.getAttributes())) {
        // existingItem and item are of incomparable specificity
        int value = valueOf(item.getConstraintValue());
        int existingValue = valueOf(existingItem.getConstraintValue());
        if (value < existingValue) {
          // item's constraint value is less than that of existingItem
          selectedItems.put(key, item);
        } else if (value == existingValue) {
          if (item.toString().compareTo(existingItem.toString()) < 0) {
            // item is ahead of existingItem in alphabetic order
            selectedItems.put(key, item);
          }
        }
      }
    }
  }
  return new HashSet<>(selectedItems.values());
}
 

private List<Message> throttle(Map<String, Integer> throttleMap, ClusterConstraints constraint,
    List<Message> messages, final boolean needThrottle) {

  List<Message> throttleOutputMsgs = new ArrayList<Message>();
  for (Message message : messages) {
    Map<ConstraintAttribute, String> msgAttr = ClusterConstraints.toConstraintAttributes(message);

    Set<ConstraintItem> matches = constraint.match(msgAttr);
    matches = selectConstraints(matches, msgAttr);

    boolean msgThrottled = false;
    for (ConstraintItem item : matches) {
      String key = item.filter(msgAttr).toString();
      if (!throttleMap.containsKey(key)) {
        throttleMap.put(key, valueOf(item.getConstraintValue()));
      }
      int value = throttleMap.get(key);
      throttleMap.put(key, --value);

      if (needThrottle && value < 0) {
        msgThrottled = true;

        if (LOG.isDebugEnabled()) {
          // TODO: printout constraint item that throttles the message
          LogUtil.logDebug(LOG, _eventId,
              "message: " + message + " is throttled by constraint: " + item);
        }
      }
    }
    if (!msgThrottled) {
      throttleOutputMsgs.add(message);
    }
  }

  return throttleOutputMsgs;
}
 

private boolean containsConstraint(Set<ConstraintItem> constraints, ConstraintItem constraint) {
  for (ConstraintItem item : constraints) {
    if (item.toString().equals(constraint.toString())) {
      return true;
    }
  }
  return false;
}
 

public ConstraintItem build() {
  // TODO: check if constraint-item is valid
  return new ConstraintItem(_attributes, _constraintValue);
}
 

@Override public void setConstraint(String clusterName,
    ClusterConstraints.ConstraintType constraintType, String constraintId,
    ConstraintItem constraintItem) {

}
 

@Test
public void testAddRemoveMsgConstraint() {
  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");

  // test admin.getMessageConstraints()
  ClusterConstraints constraints =
      tool.getConstraints(clusterName, ConstraintType.MESSAGE_CONSTRAINT);
  Assert.assertNull(constraints, "message-constraint should NOT exist for cluster: " + className);

  // remove non-exist constraint
  try {
    tool.removeConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
    // will leave a null message-constraint znode on zk
  } catch (Exception e) {
    Assert.fail("Should not throw exception when remove a non-exist constraint.");
  }

  // add a message constraint
  ConstraintItemBuilder builder = new ConstraintItemBuilder();
  builder.addConstraintAttribute(ConstraintAttribute.RESOURCE.toString(), "MyDB")
      .addConstraintAttribute(ConstraintAttribute.CONSTRAINT_VALUE.toString(), "1");
  tool.setConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1",
      builder.build());

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(clusterName);
  constraints =
      accessor.getProperty(keyBuilder.constraint(ConstraintType.MESSAGE_CONSTRAINT.toString()));
  Assert.assertNotNull(constraints, "message-constraint should exist");
  ConstraintItem item = constraints.getConstraintItem("constraint1");
  Assert.assertNotNull(item, "message-constraint for constraint1 should exist");
  Assert.assertEquals(item.getConstraintValue(), "1");
  Assert.assertEquals(item.getAttributeValue(ConstraintAttribute.RESOURCE), "MyDB");

  // test admin.getMessageConstraints()
  constraints = tool.getConstraints(clusterName, ConstraintType.MESSAGE_CONSTRAINT);
  Assert.assertNotNull(constraints, "message-constraint should exist");
  item = constraints.getConstraintItem("constraint1");
  Assert.assertNotNull(item, "message-constraint for constraint1 should exist");
  Assert.assertEquals(item.getConstraintValue(), "1");
  Assert.assertEquals(item.getAttributeValue(ConstraintAttribute.RESOURCE), "MyDB");

  // remove a exist message-constraint
  tool.removeConstraint(clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
  constraints =
      accessor.getProperty(keyBuilder.constraint(ConstraintType.MESSAGE_CONSTRAINT.toString()));
  Assert.assertNotNull(constraints, "message-constraint should exist");
  item = constraints.getConstraintItem("constraint1");
  Assert.assertNull(item, "message-constraint for constraint1 should NOT exist");

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

@Override
public void setConstraint(String clusterName, ClusterConstraints.ConstraintType constraintType, String constraintId,
    ConstraintItem constraintItem) {
  throw new IllegalStateException("Not implemented");
}
 

/**
 * Add a constraint item; create if not exist
 * @param clusterName
 * @param constraintType
 * @param constraintId
 * @param constraintItem
 */
void setConstraint(String clusterName, ConstraintType constraintType, String constraintId,
    ConstraintItem constraintItem);