@Test
public void TestSchedulerMsgContraints() throws JsonGenerationException, JsonMappingException,
IOException, InterruptedException {
TestMessagingHandlerFactoryLatch factory = new TestMessagingHandlerFactoryLatch();
HelixManager manager = null;
for (int i = 0; i < NODE_NR; i++) {
_participants[i].getMessagingService().registerMessageHandlerFactory(
factory.getMessageType(), factory);
_participants[i].getMessagingService().registerMessageHandlerFactory(
factory.getMessageType(), factory);
manager = _participants[i]; // _startCMResultMap.get(hostDest)._manager;
}
Message schedulerMessage =
new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
schedulerMessage.setTgtSessionId("*");
schedulerMessage.setTgtName("CONTROLLER");
// TODO: change it to "ADMIN" ?
schedulerMessage.setSrcName("CONTROLLER");
// Template for the individual message sent to each participant
Message msg = new Message(factory.getMessageType(), "Template");
msg.setTgtSessionId("*");
msg.setMsgState(MessageState.NEW);
// Criteria to send individual messages
Criteria cr = new Criteria();
cr.setInstanceName("localhost_%");
cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
cr.setSessionSpecific(false);
cr.setResource("%");
cr.setPartition("%");
ObjectMapper mapper = new ObjectMapper();
SerializationConfig serializationConfig = mapper.getSerializationConfig();
serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
StringWriter sw = new StringWriter();
mapper.writeValue(sw, cr);
String crString = sw.toString();
schedulerMessage.getRecord().setSimpleField("Criteria", crString);
schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");
schedulerMessage.getRecord().setSimpleField("WAIT_ALL", "true");
schedulerMessage.getRecord().setSimpleField(
DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsgContraints");
Criteria cr2 = new Criteria();
cr2.setRecipientInstanceType(InstanceType.CONTROLLER);
cr2.setInstanceName("*");
cr2.setSessionSpecific(false);
MockAsyncCallback callback = new MockAsyncCallback();
mapper = new ObjectMapper();
serializationConfig = mapper.getSerializationConfig();
serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);
sw = new StringWriter();
mapper.writeValue(sw, cr);
HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
Builder keyBuilder = helixDataAccessor.keyBuilder();
// Set contraints that only 1 msg per participant
Map<String, String> constraints = new TreeMap<String, String>();
constraints.put("MESSAGE_TYPE", "STATE_TRANSITION");
constraints.put("TRANSITION", "OFFLINE-COMPLETED");
constraints.put("CONSTRAINT_VALUE", "1");
constraints.put("INSTANCE", ".*");
manager.getClusterManagmentTool().setConstraint(manager.getClusterName(),
ConstraintType.MESSAGE_CONSTRAINT, "constraint1", new ConstraintItem(constraints));
// Send scheduler message
crString = sw.toString();
schedulerMessage.getRecord().setSimpleField("Criteria", crString);
manager.getMessagingService().sendAndWait(cr2, schedulerMessage, callback, -1);
String msgId =
callback._message.getResultMap()
.get(DefaultSchedulerMessageHandlerFactory.SCHEDULER_MSG_ID);
for (int j = 0; j < 10; j++) {
Thread.sleep(200);
PropertyKey controllerTaskStatus =
keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.toString(), msgId);
ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
if (statusUpdate.getMapFields().containsKey("SentMessageCount")) {
Assert.assertEquals(
statusUpdate.getMapFields().get("SentMessageCount").get("MessageCount"), ""
+ (_PARTITIONS * 3));
break;
}
}
for (int i = 0; i < _PARTITIONS * 3 / 5; i++) {
for (int j = 0; j < 10; j++) {
Thread.sleep(300);
if (factory._messageCount == 5 * (i + 1))
break;
}
Thread.sleep(300);
Assert.assertEquals(factory._messageCount, 5 * (i + 1));
factory.signal();
// System.err.println(i);
}
for (int j = 0; j < 10; j++) {
Thread.sleep(200);
PropertyKey controllerTaskStatus =
keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.toString(), msgId);
ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
if (statusUpdate.getMapFields().containsKey("Summary")) {
break;
}
}
Assert.assertEquals(_PARTITIONS, factory._results.size());
PropertyKey controllerTaskStatus =
keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.toString(), msgId);
ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount")
.equals("" + (_PARTITIONS * 3)));
int messageResultCount = 0;
for (String key : statusUpdate.getMapFields().keySet()) {
if (key.startsWith("MessageResult ")) {
messageResultCount++;
}
}
Assert.assertEquals(messageResultCount, _PARTITIONS * 3);
int count = 0;
for (Set<String> val : factory._results.values()) {
count += val.size();
}
Assert.assertEquals(count, _PARTITIONS * 3);
manager.getClusterManagmentTool().removeConstraint(manager.getClusterName(),
ConstraintType.MESSAGE_CONSTRAINT, "constraint1");
}